#!/usr/bin/env python3
# Coding assistance from Claude Code
"""
Dynamically select the best nearby Debian mirror by scraping the
mirror-master.debian.org status page for health scores and
cross-referencing with the Mirrors.masterlist for country and
protocol support. Falls back to a static list on failure.
"""

import argparse
import random
import sys
from html.parser import HTMLParser
from urllib.error import URLError
from urllib.request import urlopen

MIRROR_STATUS_URL = (
    'https://mirror-master.debian.org/status/mirror-status.html'
)
MASTERLIST_URL = (
    'https://mirror-master.debian.org/status/Mirrors.masterlist'
)
FETCH_TIMEOUT = 10
MAX_SCORE = 100.0
SCORE_CUTOFF = 0.90
ELIGIBLE_COUNTRIES = [
    'CA',
    'US',
]

FALLBACK_MIRRORS = [
    'debian.csail.mit.edu',
    'debian.osuosl.org',
    'mirror.cogentco.com',
    'mirror.us.leaseweb.net',
    'mirrors.xtom.com',
    'plug-mirror.rcac.purdue.edu',
]


class MirrorSelectionError(Exception):
    """Raised when live mirror data cannot be obtained or yields no candidates."""


class MirrorStatusParser(HTMLParser):
    """Parse the mirror-status HTML table to extract (site, score) pairs.

    The mirror-status page contains a single <table> whose rows look
    roughly like this (whitespace and many cells trimmed):

        <table>
          <tr>
            <th class="hostname">Site</th>
            <th>mastertrace</th>
            <th>archive version</th>
            <th>last update</th>
            <th>score</th>
            <th>...</th>
          </tr>
          <tr>
            <td data-text="debian.csail.mit.edu" class="hostname">
              <a href="http://debian.csail.mit.edu/debian/project/trace/">debian.csail.mit.edu</a>
              [<a href="mirror-hierarchy.html#debian.csail.mit.edu">H</a>,<a href="mirror-info/debian.csail.mit.edu.html">R</a>]
            </td>
            <td>...</td>          <!-- cell 2: mastertrace      -->
            <td>...</td>          <!-- cell 3: archive version  -->
            <td>...</td>          <!-- cell 4: last update      -->
            <td>100.00</td>       <!-- cell 5: score (parsed)   -->
            <td>...</td>          <!-- ...remaining cells ignored -->
          </tr>
          ...more rows, one per mirror...
        </table>

    The parser only extracts two fields per row: the hostname from the
    FIRST <a> inside cell 1 (the [H] and [R] links that follow are
    navigation and are skipped via first_link_in_cell), and the float
    in cell 5. All other cells are walked past without being inspected.
    """

    def __init__(self):
        super().__init__()
        self.in_table = False
        self.in_row = False
        self.in_cell = False
        self.in_link = False
        self.cell_index = 0
        self.current_site = None
        self.current_score = None
        self.first_link_in_cell = False
        self.mirrors = {}

    def handle_starttag(self, tag, attrs):
        # Track nesting: table > tr > td > a.
        # Each <tr> is a mirror entry; cell_index counts <td>s
        # within the row (1 = site column, 5 = score column).
        # In the site column (cell 1), the first <a> holds the
        # hostname; subsequent <a>s in that cell are navigation
        # links ([H], [R]) which we ignore via first_link_in_cell.
        if tag == 'table':
            self.in_table = True
        elif self.in_table and tag == 'tr':
            self.in_row = True
            self.cell_index = 0
            self.current_site = None
            self.current_score = None
        elif self.in_row and tag == 'td':
            self.in_cell = True
            self.cell_index += 1
            self.first_link_in_cell = True
        elif self.in_cell and tag == 'a' and self.cell_index == 1:
            if self.first_link_in_cell:
                self.in_link = True
                self.first_link_in_cell = False

    def handle_endtag(self, tag):
        # Unwind the nesting flags. When a </tr> closes, we have
        # a complete row: if both site and score were found,
        # record the entry. Rows missing either value (e.g., the
        # header row or malformed entries) are silently dropped.
        if tag == 'table':
            self.in_table = False
        elif tag == 'tr' and self.in_row:
            self.in_row = False
            if self.current_site is not None and self.current_score is not None:
                self.mirrors[self.current_site] = self.current_score
        elif tag == 'td':
            self.in_cell = False
        elif tag == 'a':
            self.in_link = False

    def handle_data(self, data):
        if self.in_link and self.cell_index == 1:
            self.current_site = data.strip()
        elif self.in_cell and self.cell_index == 5:
            text = data.strip()
            try:
                self.current_score = float(text)
            except ValueError:
                pass


class MirrorSelector:

    def __init__(self, protocol, verbose=False,
                 test_mirror_status_down=False,
                 test_master_list_down=False,
                 n_top_sites=1):
        self.protocol = protocol
        self.verbose = verbose
        self.test_mirror_status_down = test_mirror_status_down
        self.test_master_list_down = test_master_list_down
        self.n_top_sites = n_top_sites

    def log(self, msg):
        if self.verbose:
            print(f"[verbose] {msg}", file=sys.stderr)

    def fetch_url(self, url, simulate_down=False):
        if simulate_down:
            raise URLError(f"simulated outage for {url}")
        resp = urlopen(url, timeout=FETCH_TIMEOUT)
        return resp.read().decode('utf-8', errors='replace')

    def parse_mirror_scores(self, html):
        """Extract a {site: score} dict from the mirror-status HTML."""
        parser = MirrorStatusParser()
        parser.feed(html)
        return parser.mirrors

    def parse_masterlist(self, text):
        """Return {site: country_code} for eligible mirrors with protocol support."""
        if self.protocol == 'rsync':
            field = 'Archive-rsync'
        else:
            field = 'Archive-http'

        eligible = {}
        records = text.split('\n\n')
        for record in records:
            site = None
            country_code = None
            has_protocol = False
            for line in record.strip().split('\n'):
                if line.startswith('Site: '):
                    site = line[len('Site: '):].strip()
                elif line.startswith('Country: '):
                    country_code = line[len('Country: '):].split()[0]
                elif line.startswith(field + ': '):
                    has_protocol = True
            if (site is not None and country_code is not None
                    and country_code in ELIGIBLE_COUNTRIES
                    and has_protocol):
                eligible[site] = country_code

        return eligible

    def fallback(self):
        """Return a random sample of FALLBACK_MIRRORS."""
        fallbacks = random.sample(
            FALLBACK_MIRRORS,
            min(self.n_top_sites, len(FALLBACK_MIRRORS)),
        )
        self.log(f"Using fallback mirrors: {', '.join(fallbacks)}")
        return fallbacks

    def select(self):
        """Fetch live data, filter and rank mirrors, return hostname(s).

        Raises MirrorSelectionError if the data cannot be fetched, if a
        parse yields zero entries, or if no mirror passes the score
        cutoff. Callers that want the fallback behavior should catch
        the exception and call fallback() themselves.
        """
        self.log(f"Requesting {self.n_top_sites} site(s)")
        try:
            self.log(f"Fetching mirror status from {MIRROR_STATUS_URL}")
            status_html = self.fetch_url(MIRROR_STATUS_URL,
                                         self.test_mirror_status_down)
            self.log(f"Fetched mirror status ({len(status_html)} bytes)")

            self.log(f"Fetching masterlist from {MASTERLIST_URL}")
            masterlist_text = self.fetch_url(MASTERLIST_URL,
                                             self.test_master_list_down)
            self.log(f"Fetched masterlist ({len(masterlist_text)} bytes)")
        except (URLError, OSError) as e:
            raise MirrorSelectionError(
                f"could not fetch mirror data: {e}") from e

        scores = self.parse_mirror_scores(status_html)
        self.log(f"Parsed {len(scores)} mirrors from status page")
        if not scores:
            raise MirrorSelectionError(
                "mirror status page parsed to 0 entries -- "
                "the HTML format may have changed")

        eligible = self.parse_masterlist(masterlist_text)
        countries = ', '.join(ELIGIBLE_COUNTRIES)
        self.log(f"Found {len(eligible)} mirrors in [{countries}] "
                 f"with {self.protocol} support")
        if not eligible:
            raise MirrorSelectionError(
                f"masterlist yielded 0 mirrors in [{countries}] "
                f"with {self.protocol} support -- "
                "the format may have changed")

        min_score = SCORE_CUTOFF * MAX_SCORE
        self.log(f"Score cutoff: {SCORE_CUTOFF} (minimum score: {min_score})")

        candidates = []
        for site, country in eligible.items():
            if site not in scores:
                self.log(f"Skipping {site}: not found in status page")
            elif scores[site] < min_score:
                self.log(f"Skipping {site}: score {scores[site]} "
                         f"(below cutoff)")
            else:
                candidates.append((scores[site], site, country))

        if not candidates:
            raise MirrorSelectionError(
                "no suitable mirrors found "
                "(none in eligible countries met the score cutoff)")

        candidates.sort(reverse=True)
        self.log("Candidate mirrors (score, site):")
        for score, site, country in candidates:
            self.log(f"  {score:7.2f}  {site} [{country}]")

        top = candidates[:self.n_top_sites]
        selected = random.sample(
            [site for _, site, _ in top],
            len(top),
        )
        self.log(f"Randomly selected: {', '.join(selected)}")
        return selected


def main():
    parser = argparse.ArgumentParser(
        description='Select the best nearby Debian mirror')
    parser.add_argument('protocol', choices=['rsync', 'http'],
                        help='sync protocol to use')
    parser.add_argument('--verbose', action='store_true',
                        help='print debug information to stderr')
    parser.add_argument('--test-mirror-status-down', action='store_true',
                        help='simulate mirror status page being down')
    parser.add_argument('--test-master-list-down', action='store_true',
                        help='simulate master list page being down')
    parser.add_argument('--n-top-sites', type=int, default=1,
                        help='return N top sites chosen randomly')
    parser.add_argument('--use-default-on-error', action='store_true',
                        help='on any error, print mirrors from the '
                             'built-in fallback list instead of exiting '
                             'non-zero')
    args = parser.parse_args()

    selector = MirrorSelector(
        protocol=args.protocol,
        verbose=args.verbose,
        test_mirror_status_down=args.test_mirror_status_down,
        test_master_list_down=args.test_master_list_down,
        n_top_sites=args.n_top_sites,
    )
    try:
        sites = selector.select()
    except MirrorSelectionError as e:
        if args.use_default_on_error:
            print(f"Warning: {e}; using fallback", file=sys.stderr)
            sites = selector.fallback()
        else:
            print(f"Error: {e}", file=sys.stderr)
            sys.exit(1)
    for site in sites:
        print(site)


if __name__ == '__main__':
    main()
