"""Reference articles from TeachMeAnatomy and Kenhub for requests the slides do not cover.

Both sites refuse automated requests to their search pages but publish
sitemaps, so the app matches the request against the article URLs listed in
the sitemaps, downloads the best matching articles, and hands their text to
Gemini together with each page's title and URL, so the answer can name its
source exactly. Kenhub's robots.txt content signal is "use=reference": the
tutor summarizes and cites the pages, it does not reproduce them.

Pages are looked up per anatomical structure ("knee joint", "ankle joint"):
a page qualifies only through the words that name the structure, never
through a generic word such as injury, fracture, layers, or histology, so a
question about the knee cannot pull in a page about a brachial plexus injury.
"""

from __future__ import annotations

import gzip
import html
import logging
import math
import re
import threading
import time
import urllib.request
from collections import Counter, OrderedDict
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field

log = logging.getLogger("anatomy.web")

USER_AGENT = "CreightonAnatomyTutor/1.0 (educational use; fetches one article per student question)"
FETCH_TIMEOUT = 15             # seconds per download
CATALOG_TTL = 24 * 3600        # seconds before the sitemaps are downloaded again
RETRY_DELAY = 120              # seconds before a failed sitemap download is retried
MAX_PAGES = 4                  # articles handed to Gemini per request
MAX_TOPICS = 3                 # structures looked up per request
MAX_PAGE_CHARS = 12_000        # article text kept per page
PAGE_CACHE_SIZE = 64

TOKEN_RE = re.compile(r"[a-z0-9]+")
LOC_RE = re.compile(r"<loc>\s*(.*?)\s*</loc>", re.DOTALL)
TITLE_RE = re.compile(r"<title[^>]*>(.*?)</title>", re.IGNORECASE | re.DOTALL)
H1_RE = re.compile(r"<h1[^>]*>(.*?)</h1>", re.IGNORECASE | re.DOTALL)
ARTICLE_RE = re.compile(r"<article\b.*?</article>", re.IGNORECASE | re.DOTALL)
DROP_RE = re.compile(r"<(script|style|noscript|svg|nav|footer|header|form|button|iframe|figure|aside)\b.*?</\1>", re.IGNORECASE | re.DOTALL)
HEADING_RE = re.compile(r"<h([1-4])[^>]*>(.*?)</h\1>", re.IGNORECASE | re.DOTALL)
ITEM_RE = re.compile(r"<li\b[^>]*>", re.IGNORECASE)
BREAK_RE = re.compile(r"</?(p|div|br|tr|ul|ol|table|section)\b[^>]*>", re.IGNORECASE)
TAG_RE = re.compile(r"<[^>]+>", re.DOTALL)
MD_HEADING_RE = re.compile(r"(?m)^(#{1,4}) (.*)$")
CLINICAL_HEADING_RE = re.compile(r"clinic|injur|patholog|disorder|syndrome|palsy|fracture|lesion", re.IGNORECASE)


@dataclass(frozen=True)
class Site:
    name: str
    sitemap: str
    article: re.Pattern      # article URLs to keep
    children: re.Pattern | None = None   # sub-sitemaps to follow when the sitemap is an index


SITES = (
    Site(
        "TeachMeAnatomy",
        "https://teachmeanatomy.info/sitemap_index.xml",
        # Articles live three or more levels deep (/upper-limb/nerves/brachial-plexus/); two-level URLs are category
        # index pages (/upper-limb/nerves/, /encyclopaedia/d/). Skip the Spanish site and 3D models.
        re.compile(r"^https://teachmeanatomy\.info/(?!es/|3d-model/)[a-z0-9-]+/[a-z0-9-]+/[a-z0-9-]+(?:/[a-z0-9-]+)*/?$"),
        re.compile(r"/page-sitemap\d*\.xml$"),
    ),
    Site(
        "Kenhub",
        "https://www.kenhub.com/sitemap.xml.gz",
        re.compile(r"^https://www\.kenhub\.com/en/library/(?:anatomy|physiology|pathology|education)/[a-z0-9-]+$"),
    ),
)


@dataclass
class Page:
    site: str
    url: str
    title: str = ""
    text: str = ""

    @property
    def label(self) -> str:
        return f"{self.site}: {self.title}"


STOP_WORDS = set("""
a about an and any are as at be been but by can could describe did do does each explain for from give had has have
how i if in into is it its list me name of on or please should show so tell than that the their them then there
these they this those to us vs was we were what when where which while who whom whose why will with would you your
between compare define identify happen happens
""".split())
PREFIX_LENGTH = 5      # "retina" and "retinal" share the token "retin"

# Words that describe what is asked about a structure rather than naming the structure. They help rank the
# pages of a structure ("knee joint" before "knee jerk reflex") but can never make a page qualify on their
# own: "injury" alone would match "upper and lower brachial plexus injury" for a question about the knee.
GENERIC_WORDS = """
injury injuries fracture fractures dislocation dislocations sprain sprains strain strains tear tears rupture
lesion lesions damage trauma pain syndrome syndromes disease diseases disorder disorders condition conditions
clinical correlation correlations relevance significance application applications rule rules sign signs test
tests exam examination classification grade grades score method technique procedure procedures surgery
treatment management diagnosis anatomy anatomical overview introduction summary function functions structure
structures part parts region regions area areas layer layers histology system systems type types cell cells
tissue tissues blood supply innervation drainage course relation relations boundary boundaries border borders
content contents development embryology muscle muscles nerve nerves artery arteries vein veins vessel vessels
bone bones joint joints ligament ligaments tendon tendons movement movements action actions origin insertion
human body organ organs
""".split()

# The type of a structure. Generic on its own, but a page about a different type of structure with the same name
# is a different structure: "deltoid ligament" (ankle) is not "deltoid muscle" (shoulder).
TYPE_WORDS = """
muscle muscles nerve nerves artery arteries vein veins ligament ligaments tendon tendons bone bones joint joints
ganglion bursa cartilage duct gland fascia membrane
""".split()

# Technical joint names, and the common names the articles are filed under.
JOINT_NAMES = {
    "talocrural": "ankle", "glenohumeral": "shoulder", "tibiofemoral": "knee", "femorotibial": "knee",
    "acetabulofemoral": "hip", "coxofemoral": "hip", "radiocarpal": "wrist", "humeroulnar": "elbow",
    "humeroradial": "elbow",
}

# Names the slides use for structures that both sites file under another name.
SYNONYMS = {"calcaneal tendon": "achilles tendon", "tendo calcaneus": "achilles tendon"}

# Position words narrow a name but are often left out of article titles ("lateral ankle ligaments" is covered by
# the ankle joint article), so a page may lack them; every other naming word must appear in the page's URL.
POSITION_WORDS = """
anterior posterior superior inferior medial lateral proximal distal deep superficial internal external left right
upper lower common greater lesser major minor long short
""".split()


def tokenize(text: str) -> list[str]:
    """Lower-case words without stop words (URL slugs are full of "of" and "the"), trailing plural s dropped,
    plus a five-letter prefix for longer words so that inflections such as retina/retinal still match."""
    tokens = []
    for tok in TOKEN_RE.findall(text.lower()):
        if tok in STOP_WORDS:
            continue
        if len(tok) > 3 and tok.endswith("s"):
            tok = tok[:-1]
        tokens.append(tok)
        if len(tok) > PREFIX_LENGTH:
            tokens.append(tok[:PREFIX_LENGTH])
    return tokens


class BM25:
    """Small BM25 index over the URL slugs."""

    def __init__(self, documents: list[list[str]], k1: float = 1.5, b: float = 0.75):
        self.k1, self.b = k1, b
        self.doc_len = [len(d) for d in documents]
        self.avg_len = (sum(self.doc_len) / len(documents)) if documents else 1.0
        self.postings: dict[str, list[tuple[int, int]]] = {}
        for idx, doc in enumerate(documents):
            for term, freq in Counter(doc).items():
                self.postings.setdefault(term, []).append((idx, freq))
        n = len(documents)
        self.idf = {t: math.log(1 + (n - len(p) + 0.5) / (len(p) + 0.5)) for t, p in self.postings.items()}

    def scores(self, query: list[str]) -> list[float]:
        out = [0.0] * len(self.doc_len)
        for term in set(query):
            idf = self.idf.get(term)
            if idf is None:
                continue
            for idx, freq in self.postings[term]:
                denom = freq + self.k1 * (1 - self.b + self.b * self.doc_len[idx] / self.avg_len)
                out[idx] += idf * freq * (self.k1 + 1) / denom
        return out


def fetch(url: str, timeout: int = FETCH_TIMEOUT) -> bytes:
    request = urllib.request.Request(
        url, headers={"User-Agent": USER_AGENT, "Accept": "text/html,application/xml;q=0.9,*/*;q=0.8", "Accept-Language": "en"}
    )
    with urllib.request.urlopen(request, timeout=timeout) as response:
        data = response.read()
    if data[:2] == b"\x1f\x8b":
        data = gzip.decompress(data)
    return data


def sitemap_urls(site: Site) -> list[str]:
    """Article URLs listed in a site's sitemap (one level of sub-sitemaps is followed)."""
    xml = fetch(site.sitemap).decode("utf-8", "ignore")
    locations = [html.unescape(u) for u in LOC_RE.findall(xml)]
    if "<sitemapindex" in xml:
        urls: list[str] = []
        for child in locations:
            if site.children is None or site.children.search(child):
                urls += [html.unescape(u) for u in LOC_RE.findall(fetch(child).decode("utf-8", "ignore"))]
        locations = urls
    return [u for u in locations if site.article.match(u)]


def url_tokens(url: str) -> list[str]:
    path = url.split("://", 1)[-1].split("/", 1)[-1]
    return tokenize(path.replace("-", " ").replace("/", " "))


def url_path(url: str) -> str:
    return url.split("://", 1)[-1].split("/", 1)[-1].replace("-", " ").replace("/", " ")


def singular(tok: str) -> str:
    """"joints" -> "joint", "arches" -> "arch", "sinuses" -> "sinus", Latin "menisci" -> "meniscus",
    "vertebrae" -> "vertebra"; words ending in -us, -is, or -ss are left alone ("trapezius", "pelvis")."""
    if len(tok) > 4 and tok.endswith("ae"):
        return tok[:-1]
    if len(tok) > 4 and tok.endswith("i"):
        return tok[:-1] + "us"
    if len(tok) > 4 and tok.endswith(("ches", "shes", "xes", "zes", "uses", "sses")):
        return tok[:-2]
    if len(tok) > 3 and tok.endswith("s") and not tok.endswith(("us", "is", "ss")):
        return tok[:-1]
    return tok


def words(text: str) -> list[str]:
    """Whole words of a text, stop words dropped, singular, technical joint names replaced by common ones."""
    out = []
    for tok in TOKEN_RE.findall(text.lower()):
        if tok in STOP_WORDS:
            continue
        tok = singular(JOINT_NAMES.get(tok, tok))
        out.append(tok)
    return out


def same_word(a: str, b: str) -> bool:
    """Equal, or the same stem with a short ending: "retina"/"retinal", "cochlea"/"cochlear"."""
    if a == b:
        return True
    short, long_ = sorted((a, b), key=len)
    return len(short) >= 5 and len(long_) - len(short) <= 2 and long_.startswith(short)


GENERIC = frozenset(w for text in GENERIC_WORDS for w in words(text))
POSITION = frozenset(w for text in POSITION_WORDS for w in words(text))
TYPES = frozenset(w for text in TYPE_WORDS for w in words(text))


@dataclass
class Catalog:
    """Article URLs of both sites with a BM25 index over their slugs, refreshed daily."""

    # ((site name, url) entries, BM25 index over their slugs, whole words of each slug): replaced as one object
    data: tuple[list[tuple[str, str]], BM25 | None, list[frozenset]] = ([], None, [])
    next_load: float = 0.0
    error: str = ""
    lock: threading.Lock = field(default_factory=threading.Lock)

    def ensure(self) -> str:
        """Load or refresh the catalog when due. Returns a note about sites that could not be reached."""
        with self.lock:
            now = time.monotonic()
            if now < self.next_load:
                return self.error
            self.next_load = now + RETRY_DELAY
            entries: list[tuple[str, str]] = []
            errors: list[str] = []
            for site in SITES:
                try:
                    urls = sitemap_urls(site)
                    log.info("%s: %d articles listed", site.name, len(urls))
                    entries += [(site.name, url) for url in urls]
                except Exception as err:
                    log.warning("%s sitemap failed: %s", site.name, err)
                    errors.append(f"{site.name} could not be reached")
            if entries:
                self.data = (
                    entries,
                    BM25([url_tokens(url) for _, url in entries]),
                    [frozenset(words(url_path(url))) for _, url in entries],
                )
                if not errors:
                    self.next_load = now + CATALOG_TTL
            self.error = "; ".join(errors)
            return self.error


CATALOG = Catalog()
_page_cache: OrderedDict[str, Page] = OrderedDict()
_page_lock = threading.Lock()


def warm_up() -> None:
    """Download the sitemaps in the background so the first fallback request is fast."""
    threading.Thread(target=CATALOG.ensure, name="web-catalog", daemon=True).start()


def article_text(page_html: str) -> str:
    """Plain text of an article page: the <article> element, headings and list items kept as Markdown."""
    match = ARTICLE_RE.search(page_html)
    body = match.group(0) if match else page_html
    body = DROP_RE.sub(" ", body)
    # Heading text is put on one line: Kenhub headings contain line breaks, which left "##" alone on a line.
    body = HEADING_RE.sub(
        lambda m: "\n\n" + "#" * int(m.group(1)) + " " + " ".join(TAG_RE.sub(" ", m.group(2)).split()) + "\n", body)
    body = ITEM_RE.sub("\n- ", body)
    body = BREAK_RE.sub("\n", body)
    text = html.unescape(TAG_RE.sub(" ", body))
    text = re.sub(r"[ \t\r\f\v]+", " ", text)
    text = re.sub(r"(?m)^ +| +$", "", text)
    text = re.sub(r"(?m)^-\s*$\n?", "", text)          # empty list items (media controls)
    text = re.sub(r"\n\s*\n+", "\n\n", text)
    return text.strip()


def excerpt(text: str, limit: int = MAX_PAGE_CHARS) -> str:
    """The first `limit` characters of an article, except that clinical sections ("Clinical significance",
    "Brachial plexus injury") are kept even when they come later: Kenhub puts them at the end of long articles.
    Clinical sections get at most half of the budget."""
    if len(text) <= limit:
        return text
    headings = [(m.start(), len(m.group(1)), m.group(2)) for m in MD_HEADING_RE.finditer(text)]
    sections: list[tuple[int, int]] = []
    for i, (start, level, title) in enumerate(headings):
        if not CLINICAL_HEADING_RE.search(title) or (sections and start < sections[-1][1]):
            continue
        end = next((s for s, lvl, _ in headings[i + 1:] if lvl <= level), len(text))
        sections.append((start, end))
    head_end = limit
    tail = ""
    for start, end in sections:
        if start < limit // 2:            # already inside the part that is kept
            continue
        piece = text[start:end].strip()
        if len(tail) + len(piece) > limit // 2:
            piece = piece[: max(0, limit // 2 - len(tail))]
        tail = (tail + "\n\n" + piece).strip() if piece else tail
        if len(tail) >= limit // 2:
            break
    if tail:
        head_end = limit - len(tail) - 9
    return (text[:head_end].rstrip() + ("\n\n[...]\n\n" + tail if tail else "")).strip()


def page_title(page_html: str, site: str) -> str:
    match = H1_RE.search(page_html) or TITLE_RE.search(page_html)
    title = html.unescape(TAG_RE.sub("", match.group(1))).strip() if match else ""
    title = re.sub(r"\s*\|\s*(Kenhub|TeachMeAnatomy).*$", "", title, flags=re.IGNORECASE)
    return " ".join(title.split()) or site


def load_page(page: Page) -> Page | None:
    """Download one article (cached). Returns None when it cannot be read."""
    with _page_lock:
        cached = _page_cache.get(page.url)
        if cached is not None:
            _page_cache.move_to_end(page.url)
            return cached
    try:
        page_html = fetch(page.url).decode("utf-8", "ignore")
    except Exception as err:
        log.warning("Could not fetch %s: %s", page.url, err)
        return None
    page.title = page_title(page_html, page.site)
    page.text = excerpt(article_text(page_html))
    if not page.text:
        return None
    with _page_lock:
        _page_cache[page.url] = page
        while len(_page_cache) > PAGE_CACHE_SIZE:
            _page_cache.popitem(last=False)
    return page


def topic_matches(topic: str, entries: list[tuple[str, str]], index: BM25, slug_words: list[frozenset]) -> list[int]:
    """Catalog entries about one structure, best first.

    The naming words of the phrase are its words that are not generic ("knee joint" -> knee; "tarsal tunnel" ->
    tarsal, tunnel). A page qualifies only if its URL contains every naming word except position words such as
    lateral or superior, so "tarsal tunnel" never qualifies "carpal tunnel" and "superior thoracic aperture" never
    qualifies "superior thoracic artery". When the phrase names the type of structure ("deltoid ligament"), a page
    whose URL names another type ("deltoid muscle") does not qualify. Pages that contain more naming words come
    first, then pages that name the same type, then pages whose URL has fewer other words ("retina" before
    "central retinal artery"), then the BM25 score of the whole phrase.
    """
    phrase_words = words(topic)
    naming = list(dict.fromkeys(w for w in phrase_words if w not in GENERIC))
    required = [w for w in naming if w not in POSITION] or naming[-1:]
    if not required:
        return []
    phrase = set(phrase_words)
    types = phrase & TYPES

    def has(word: str, slug: frozenset) -> bool:
        return any(same_word(word, other) for other in slug)

    found = [i for i in range(len(entries)) if all(has(w, slug_words[i]) for w in required)]
    if types:
        found = [i for i in found if not (slug_words[i] & TYPES) or slug_words[i] & types]
    if not found:
        return []
    hits = {i: sum(1 for w in naming if has(w, slug_words[i])) for i in found}
    # Naming words in the article's own slug, not only in its category path ("pelvis/bones/sacrum" is filed under
    # pelvis but is about the sacrum).
    own = {i: sum(1 for w in naming if has(w, frozenset(words(url_path(entries[i][1].rstrip("/").rsplit("/", 1)[-1])))))
           for i in found}
    typed = {i: bool(slug_words[i] & types) for i in found}
    extra = {i: sum(1 for w in slug_words[i] if not any(same_word(w, q) for q in phrase)) for i in found}
    scores = index.scores(tokenize(topic))
    return sorted(found, key=lambda i: (-hits[i], -own[i], not typed[i], extra[i], -scores[i], len(entries[i][1])))


def names_structure(topic: str, page: Page) -> bool:
    """For a phrase that names the type of structure ("deltoid ligament"), the page must mention that structure:
    the name followed within a few words by the type. Other phrases are always accepted."""
    raw = [t for t in TOKEN_RE.findall(topic.lower()) if t not in STOP_WORDS]
    type_words = [t for t in raw if singular(t) in TYPES]
    names = [t for t in raw if singular(JOINT_NAMES.get(t, t)) not in GENERIC | TYPES]
    if not type_words or not names:
        return True
    name = JOINT_NAMES.get(names[-1], names[-1])
    stem = singular(name)[:max(4, len(singular(name)) - 2)]
    kind = singular(type_words[0])
    pattern = re.compile(rf"\b{re.escape(stem)}\w*\W+(?:\w+\W+){{0,3}}{re.escape(kind)}", re.IGNORECASE)
    return bool(pattern.search(page.title + "\n" + page.text))


def question_topics(question: str) -> list[str]:
    """Structures to look up when the model named none: the question as one phrase, or, when no article matches
    the whole phrase ("ligaments of the knee and ankle"), each of its naming words ("knee", "ankle")."""
    CATALOG.ensure()
    entries, index, slug_words = CATALOG.data
    if index is None or topic_matches(question, entries, index, slug_words):
        return [question]
    naming = [w for w in dict.fromkeys(words(question)) if w not in GENERIC and w not in POSITION]
    return naming[:MAX_TOPICS] or [question]


def find_pages(topics: list[str] | str) -> tuple[list[Page], str]:
    """Articles about the given structures, text included, plus a note about sites that could not be reached.

    Each structure gets at most one page per site; pages are taken structure by structure, best first, so a
    request about the knee and the ankle gets pages about both.
    """
    if isinstance(topics, str):
        topics = [topics]
    topics = [t.strip() for t in topics if t and t.strip()][:MAX_TOPICS]
    for name, other in SYNONYMS.items():
        topics = [re.sub(rf"\b{name}\b", other, t, flags=re.IGNORECASE) for t in topics]
    error = CATALOG.ensure()
    entries, index, slug_words = CATALOG.data
    if index is None or not topics:
        return [], error
    per_topic: list[list[tuple[str, Page]]] = []
    for topic in topics:
        picks: dict[str, Page] = {}
        for idx in topic_matches(topic, entries, index, slug_words):
            site, url = entries[idx]
            picks.setdefault(site, Page(site, url))
        per_topic.append([(topic, page) for page in picks.values()])      # best page first
    chosen: list[tuple[str, Page]] = []
    seen: set[str] = set()
    for rank in range(len(SITES)):                     # every structure's best page first, then the second site
        for picks in per_topic:
            if rank < len(picks) and picks[rank][1].url not in seen and len(chosen) < MAX_PAGES:
                seen.add(picks[rank][1].url)
                chosen.append(picks[rank])
    if not chosen:
        return [], error
    with ThreadPoolExecutor(max_workers=len(chosen)) as pool:
        loaded = list(pool.map(load_page, [page for _, page in chosen]))
    pages = [page for (topic, _), page in zip(chosen, loaded) if page is not None and names_structure(topic, page)]
    if not pages and any(page is None for page in loaded):
        error = "; ".join(filter(None, [error, "the matching TeachMeAnatomy and Kenhub pages could not be downloaded"]))
    return pages, error