"""
retrieve.py — Hybrid retrieval: semantic search + title keyword matching
              + live Necyklopedie fallback for articles not in local index.
"""

import json
import os
import re
import time

from json_storage import JsonValue, update_json

_title_index = None

# ── Live-fetch rate limiting ───────────────────────────────────────────────────
_LIVE_LOG = os.path.join("data", "live_fetch_log.json")
_LIVE_MAX_DAY = 300  # max live fetches per 24 h
_LIVE_MAX_MIN = 30   # max live fetches per 60 s
_CHUNK_SIZE   = 800  # chars per chunk for live-fetched text


def _load_title_index():
    """Build a lowercase title -> pageid lookup from articles.jsonl."""
    global _title_index
    if _title_index is not None:
        return _title_index

    _title_index = {}
    articles_path = os.path.join("data", "articles.jsonl")
    if os.path.exists(articles_path):
        with open(articles_path, "r", encoding="utf-8") as f:
            for line in f:
                rec = json.loads(line)
                _title_index[rec["title"].lower()] = rec["pageid"]
    return _title_index


def _czech_stem(word):
    """Crude Czech stemming: strip common suffixes for matching."""
    # Order matters: try longer suffixes first
    for suffix in ["ách", "ům", "emi", "ami", "ích", "ech",
                    "ám", "ím", "ov", "ěm",
                    "ou", "em", "mi", "ně",
                    "ě", "y", "u", "i", "a", "o"]:
        if word.endswith(suffix) and len(word) - len(suffix) >= 3:
            return word[:-len(suffix)]
    return word


def _stems_match(query_stem, title_stem):
    """Check if query stem matches title stem. Allows 1 char difference for Czech declension."""
    if query_stem == title_stem:
        return True
    if len(query_stem) >= 3 and title_stem.startswith(query_stem) and len(title_stem) - len(query_stem) == 1:
        return True
    return len(title_stem) >= 3 and query_stem.startswith(title_stem) and len(query_stem) - len(title_stem) == 1


def _find_matching_titles(query_lower):
    """Find article titles that match query words. Returns dict of pageid -> match_quality.
    Match quality: 3 = exact title match, 2 = title is a query word, 1 = substring match."""
    idx = _load_title_index()
    stopwords = {
        "se", "si", "to", "je", "na", "ve", "za", "ke", "od", "do", "po", "při",
        "pro", "nad", "pod", "jak", "kdo", "kde", "kdy", "cos", "ten", "the",
        "jeho", "její", "jsou", "bylo", "být", "mít", "které", "která", "který",
        "jako", "ale", "ani", "nebo", "tak", "jen", "ještě", "což", "tedy",
        "taky", "také", "toho", "tím", "tom", "této", "tohle", "nic", "než",
        "aby", "když", "před", "mezi", "přes", "bez", "stál", "řekl", "jsem",
        "jsi", "jsme", "jste", "byl", "byla", "byli", "není", "nás", "vás",
        "mně", "mne", "jím", "jich", "jim", "ním", "nich", "něm", "nej",
        "své", "svou", "svým", "jaký", "jaká", "jaké", "celý", "celá",
        "dnes", "bude", "více", "méně", "setkal", "říkají", "pravda",
        "pojem", "pojmem", "pojmu", "vysvětli", "vysvětlit", "řekni", "popiš", "popsat",
        "říká", "neříká", "proboha", "skrývá", "město", "města",
        "vůbec", "nevím", "utajované", "informace", "pouč", "slyšel", "nikdy", "neslyšel", "vzniklo", "vznikl", "vznikla",
    }
    raw_words = re.findall(r'[\w]+', query_lower)  # strip punctuation
    words = [w for w in raw_words if len(w) >= 3 and w not in stopwords]
    stems = [_czech_stem(w) for w in words]

    matched = {}  # pageid -> match_quality
    query_joined = " ".join(words)

    for title_lower, pageid in idx.items():
        title_stems = [_czech_stem(tw) for tw in title_lower.split()]
        title_stem = _czech_stem(title_lower)
        best = 0
        # Check if full query matches full title (multi-word exact match)
        # Require title >= 5 chars for substring match to avoid short titles ("Pod", "Sad"…)
        # matching inside unrelated long words ("blaosdpsadkpods")
        if title_lower == query_joined or (len(title_lower) >= 5 and title_lower in query_joined):
            best = 3
        else:
            for word, stem in zip(words, stems):
                # Exact single-word title match
                if word == title_lower or stem == title_stem:
                    best = max(best, 3)
                # Query word matches a title word (multi-word title)
                elif word in title_lower.split() or stem in title_stems:
                    best = max(best, 2)
                # Substring match (weakest)
                elif word in title_lower:
                    best = max(best, 1)
        if best > 0:
            matched[pageid] = best

    return matched


# ── Live Necyklopedie fallback ─────────────────────────────────────────────────

def _reserve_live_fetch() -> bool:
    """Atomically reserve one live fetch when both quotas allow it."""
    now = time.time()
    is_allowed = False

    def reserve(log: list[JsonValue]) -> None:
        nonlocal is_allowed
        recent: list[int | float] = []
        for timestamp in log:
            if isinstance(timestamp, (int, float)) and now - timestamp < 86400:
                recent.append(timestamp)
        log[:] = recent
        if len(recent) >= _LIVE_MAX_DAY:
            return
        if sum(1 for timestamp in recent if now - timestamp < 60) >= _LIVE_MAX_MIN:
            return
        log.append(now)
        is_allowed = True

    try:
        update_json(_LIVE_LOG, list, reserve)
    except OSError:
        return False
    return is_allowed


def _query_words(query_lower):
    """Return content words from query (same stopwords as title matching)."""
    stopwords = {
        "se", "si", "to", "je", "na", "ve", "za", "ke", "od", "do", "po", "při",
        "pro", "nad", "pod", "jak", "kdo", "kde", "kdy", "cos", "ten", "the",
        "jeho", "její", "jsou", "bylo", "být", "mít", "které", "která", "který",
        "jako", "ale", "ani", "nebo", "tak", "jen", "ještě", "což", "tedy",
        "taky", "také", "toho", "tím", "tom", "této", "tohle", "nic", "než",
        "aby", "když", "před", "mezi", "přes", "bez", "stál", "řekl", "jsem",
        "jsi", "jsme", "jste", "byl", "byla", "byli", "není", "nás", "vás",
        "mně", "mne", "jím", "jich", "jim", "ním", "nich", "něm", "nej",
        "své", "svou", "svým", "jaký", "jaká", "jaké", "celý", "celá",
        "dnes", "bude", "více", "méně", "setkal", "říkají", "pravda",
        "pojem", "pojmem", "pojmu", "vysvětli", "vysvětlit", "řekni", "popiš", "popsat",
        "říká", "neříká", "proboha", "skrývá", "město", "města",
        "vůbec", "nevím", "utajované", "informace", "pouč", "slyšel", "nikdy", "neslyšel", "vzniklo", "vznikl", "vznikla",
        # English
        "what", "who", "how", "where", "when", "why", "tell", "about", "explain",
        "is", "are", "was", "were", "this", "that", "with", "from",
    }
    return [w for w in re.findall(r'[\w]+', query_lower) if len(w) >= 3 and w not in stopwords]


_API_URL = "https://necyklopedie.org/w/api.php"
_UA = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36"


def _strip_diacritics(s):
    import unicodedata
    return "".join(c for c in unicodedata.normalize("NFD", s)
                   if unicodedata.category(c) != "Mn")


def _search_necyklopedie(word):
    """Find Necyklopedie article titles matching word, handling diacritics.

    Strategy: allpages prefix search with progressively shorter prefixes of the
    diacritics-stripped word. MediaWiki prefix search is Unicode-aware at the
    char level — 'Podprd' matches 'Podprdélník' but 'Podprde' doesn't (é≠e).
    We then keep candidates whose stripped title closely matches our stripped word.
    """
    import urllib.parse
    import urllib.request

    stripped = _strip_diacritics(word.lower())
    target_stem = stripped[:max(4, len(stripped) - 2)]  # at least 4 chars

    # Try from full length down to 4 chars, stop at first hit
    for length in range(len(stripped), 3, -1):
        prefix = stripped[:length].capitalize()
        params = urllib.parse.urlencode({
            "action": "query", "list": "allpages",
            "apprefix": prefix, "apnamespace": "0",
            "aplimit": "10", "format": "json",
        })
        url = f"{_API_URL}?{params}"
        req = urllib.request.Request(url, headers={"User-Agent": _UA})
        with urllib.request.urlopen(req, timeout=10) as resp:
            data = json.loads(resp.read(_MAX_SEARCH_BYTES).decode("utf-8"))
        pages = data.get("query", {}).get("allpages", [])
        if not pages:
            continue
        # Keep candidates whose stripped title starts with our stripped word stem
        matches = [
            p["title"] for p in pages
            if _strip_diacritics(p["title"].lower()).startswith(target_stem)
        ]
        if matches:
            return matches[:3]
        # If pages found but none match stem, stop — we've gone too broad
        if len(pages) >= 5:
            break

    return []


_MAX_EXPORT_BYTES = 4 * 1024 * 1024
_MAX_SEARCH_BYTES = 256 * 1024


def _fetch_pages_by_title(titles):
    """Fetch specific pages through the API: {pageid: {pageid, title, text}}.
    The reply is read inside a chat request, so it is capped (ValueError)."""
    import scrape
    return scrape.fetch_pages(titles, timeout=15, max_bytes=_MAX_EXPORT_BYTES)


def _live_fetch(query_lower):
    """Fetch article from Necyklopedie live. Returns list of (doc, meta) chunks or []."""
    words = _query_words(query_lower)
    if not words or not _reserve_live_fetch():
        return []

    # Use longest word as primary search term (most specific)
    primary = max(set(words), key=len)

    # Use MediaWiki search to find the real title (handles diacritics, capitalization)
    try:
        titles_to_try = _search_necyklopedie(primary)
    except Exception:  # noqa: BLE001 - best effort (network, HTTP, JSON); fall back to a guessed title
        titles_to_try = [primary.capitalize()]

    if not titles_to_try:
        return []

    try:
        results = _fetch_pages_by_title(titles_to_try[:3])
    except Exception:  # noqa: BLE001 - best effort (network, HTTP, XML); no live chunks
        return []

    if not results:
        return []

    chunks = []
    for pageid, rec in results.items():
        text = rec["text"]
        title = rec["title"]
        # Split into fixed-size chunks
        for i in range(0, min(len(text), _CHUNK_SIZE * 5), _CHUNK_SIZE):
            chunk = text[i:i + _CHUNK_SIZE]
            if chunk.strip():
                chunks.append((chunk, {
                    "title": title,
                    "pageid": pageid,
                    "chunk_index": i // _CHUNK_SIZE,
                    "live": True,
                }))

    return chunks


def retrieve_chunks(query, embedder, collection, top_k=10):
    """Retrieve relevant chunks using semantic search + title keyword boost.

    Returns (chunks, live_titles, live_attempted) where live_titles is a list of
    article titles that were fetched live from Necyklopedie (empty if all from
    local index) and live_attempted says whether a live fetch was tried.
    """
    query_lower = query.lower()

    # 1. Semantic search
    query_embedding = embedder.encode(query).tolist()
    sem_results = collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k,
    )
    sem_chunks = list(zip(
        sem_results["documents"][0],
        sem_results["metadatas"][0],
    ))

    # 2. Title keyword matching via local index
    matched = _find_matching_titles(query_lower)  # pageid -> match_quality

    # Limit to top 5 articles by match quality to avoid excessive DB calls
    top_matched = sorted(matched.items(), key=lambda x: -x[1])[:5]

    article_chunks = []  # list of (match_quality, [(doc, meta), ...])
    for pageid, match_quality in top_matched:
        try:
            kw_results = collection.get(
                where={"pageid": pageid},
                include=["documents", "metadatas"],
            )
            if not kw_results["documents"]:
                continue
            pairs = list(zip(kw_results["documents"], kw_results["metadatas"]))
            pairs.sort(key=lambda p: p[1].get("chunk_index", 0))
            picks = pairs[:5]
            article_chunks.append((match_quality, picks))
        except Exception:  # noqa: BLE001, S110 - one failed lookup must not drop the other articles
            pass

    article_chunks.sort(reverse=True, key=lambda x: x[0])
    title_chunks = []
    for _, chunks in article_chunks:
        title_chunks.extend(chunks)

    # 3. Merge: title matches first (deduplicated), then semantic results
    seen_ids = set()
    merged = []
    for source in [title_chunks, sem_chunks]:
        for doc, meta in source:
            chunk_id = f"{meta['pageid']}_{meta['chunk_index']}"
            if chunk_id not in seen_ids:
                seen_ids.add(chunk_id)
                merged.append((doc, meta))

    # 4. Live fallback: if no keyword title match at all, ask Necyklopedie directly
    live_titles = []
    live_attempted = False
    if not matched:
        live_attempted = True
        live_chunks = _live_fetch(query_lower)
        if live_chunks:
            live_titles = list(dict.fromkeys(m["title"] for _, m in live_chunks))
            # Prepend live chunks (they're the best match we have)
            seen_live = set()
            deduped_live = []
            for doc, meta in live_chunks:
                cid = f"{meta['pageid']}_{meta['chunk_index']}"
                if cid not in seen_ids and cid not in seen_live:
                    seen_live.add(cid)
                    deduped_live.append((doc, meta))
            merged = deduped_live + merged

    return merged[:top_k], live_titles, live_attempted