"""
Fusion + ranking.

Pipeline:
  1. Merge every provider's SERP into unique Documents (canonical-URL dedup).
  2. Reciprocal Rank Fusion across provider rankings (weighted by trust).
  3. BM25 lexical relevance against the query (+ expansions).
  4. Authority prior, freshness decay, cross-provider agreement.
  5. Optional LLM reranking of the top slice (a cheap cross-encoder stand-in).
  6. Domain diversity enforcement (MMR-ish) so one site can't dominate.
"""
from __future__ import annotations

import math
import re
from collections import defaultdict
from datetime import datetime, timezone
from typing import Dict, Iterable, List, Optional, Sequence, Tuple

from rank_bm25 import BM25Okapi

from .config import AUTHORITY, DOMAIN_PENALTIES, settings
from .schemas import Chunk, Document, RawResult, domain_of, registrable

_TOKEN = re.compile(r"[a-z0-9]+")
_STOP = {
    "the", "a", "an", "and", "or", "of", "to", "in", "is", "are", "was", "were",
    "for", "on", "with", "as", "by", "at", "from", "that", "this", "it", "be",
    "how", "what", "why", "when", "where", "which", "who", "does", "do", "did",
    "can", "could", "should", "would", "will", "i", "you", "we", "they", "my",
}


def tokenize(text: str, keep_stop: bool = False) -> List[str]:
    toks = _TOKEN.findall((text or "").lower())
    return toks if keep_stop else [t for t in toks if t not in _STOP and len(t) > 1]


# --------------------------------------------------------------------------
# 1. merge
# --------------------------------------------------------------------------
def merge_results(
    per_query: Dict[str, List[RawResult]],
) -> List[Document]:
    """Collapse provider rows into unique Documents keyed by canonical URL."""
    docs: Dict[str, Document] = {}
    for query, rows in per_query.items():
        for r in rows:
            d = docs.get(r.url)
            if d is None:
                d = Document(
                    url=r.url, title=r.title, snippet=r.snippet,
                    domain=domain_of(r.url), published=r.published,
                    extra=dict(r.extra or {}),
                )
                docs[r.url] = d
            # keep the richest metadata we've seen
            if len(r.title) > len(d.title):
                d.title = r.title
            if len(r.snippet) > len(d.snippet):
                d.snippet = r.snippet
            if r.published and not d.published:
                d.published = r.published
            if r.extra:
                for k, v in r.extra.items():
                    d.extra.setdefault(k, v)
            d.providers.append(r.provider)
            prev = d.ranks.get(r.provider)
            if prev is None or r.rank < prev:
                d.ranks[r.provider] = r.rank
            if query not in d.queries:
                d.queries.append(query)
    return list(docs.values())


# --------------------------------------------------------------------------
# 2-4. signals
# --------------------------------------------------------------------------
def _provider_weight(name: str) -> float:
    from .providers import PROVIDERS_BY_NAME
    p = PROVIDERS_BY_NAME.get(name.split(":")[0])
    return p.weight if p else 1.0


def compute_rrf(docs: Sequence[Document], k: int = 60) -> None:
    for d in docs:
        score = 0.0
        for provider, rank in d.ranks.items():
            score += _provider_weight(provider) / (k + rank)
        # small bonus for appearing under multiple distinct sub-queries
        score *= 1.0 + 0.12 * (len(d.queries) - 1)
        d.rrf = score
    _normalize(docs, "rrf")


def compute_bm25(docs: Sequence[Document], queries: Sequence[str]) -> None:
    if not docs:
        return
    corpus = []
    for d in docs:
        # title is worth more than body: repeat it
        blob = f"{d.title} {d.title} {d.domain} {d.snippet} {d.content[:4000]}"
        corpus.append(tokenize(blob))
    if not any(corpus):
        return
    try:
        bm = BM25Okapi(corpus)
    except Exception:
        return
    agg = [0.0] * len(docs)
    for qi, q in enumerate(queries):
        qt = tokenize(q)
        if not qt:
            continue
        weight = 1.0 if qi == 0 else 0.65  # original query dominates
        scores = bm.get_scores(qt)
        for i, s in enumerate(scores):
            agg[i] += float(s) * weight
    for d, s in zip(docs, agg):
        d.bm25 = s
    _normalize(docs, "bm25")


def compute_authority(docs: Iterable[Document]) -> None:
    for d in docs:
        # A magnet link has no host to judge. Trust it on swarm health
        # instead, otherwise every file result sinks below any web page.
        if d.url.startswith("magnet:"):
            seeders = 0
            try:
                seeders = int((d.extra or {}).get("seeders") or 0)
            except (TypeError, ValueError):
                seeders = 0
            d.authority = max(0.55, min(0.95, 0.55 + seeders / 2000.0))
            continue
        reg = registrable(d.domain)
        base = AUTHORITY.get(d.domain) or AUTHORITY.get(reg)
        if base is None:
            base = 0.5
            if reg.endswith((".gov", ".edu", ".mil")) or reg.endswith((".gov.in", ".ac.uk", ".edu.au")):
                base = 0.88
            elif reg.endswith(".org"):
                base = 0.60
            elif reg.endswith((".io", ".dev", ".ai")):
                base = 0.55
            elif re.search(r"\d{4,}", reg) or reg.count("-") >= 3:
                base = 0.32  # spammy-looking host
        penalty = DOMAIN_PENALTIES.get(d.domain, DOMAIN_PENALTIES.get(reg, 0.0))
        d.authority = max(0.0, min(1.0, base + penalty))


_DATE_PATTERNS = (
    "%Y-%m-%dT%H:%M:%S%z", "%Y-%m-%dT%H:%M:%SZ", "%Y-%m-%dT%H:%M:%S",
    "%Y-%m-%d %H:%M:%S", "%Y-%m-%d", "%Y/%m/%d", "%d %b %Y", "%b %d, %Y",
    "%Y %b %d", "%Y",
)


def parse_date(value: Optional[str]) -> Optional[datetime]:
    if not value:
        return None
    v = str(value).strip()
    v = re.sub(r"\.\d+Z?$", "", v)
    for fmt in _DATE_PATTERNS:
        try:
            dt = datetime.strptime(v[:32], fmt)
            return dt.replace(tzinfo=dt.tzinfo or timezone.utc)
        except ValueError:
            continue
    m = re.search(r"(20[0-2]\d|19\d\d)", v)
    if m:
        try:
            return datetime(int(m.group(1)), 6, 15, tzinfo=timezone.utc)
        except ValueError:
            return None
    return None


def compute_freshness(docs: Iterable[Document], half_life_days: float = 540.0) -> None:
    now = datetime.now(timezone.utc)
    for d in docs:
        dt = parse_date(d.published)
        if not dt:
            d.freshness = 0.45  # unknown: neutral-ish
            continue
        age = max(0.0, (now - dt).total_seconds() / 86400.0)
        d.freshness = math.exp(-age / half_life_days)


def compute_agreement(docs: Sequence[Document]) -> None:
    """
    Cross-source corroboration: documents whose distinctive terms recur
    across *other domains* are more likely to be factually grounded.
    """
    if len(docs) < 2:
        for d in docs:
            d.agreement = 0.5
        return
    df: Dict[str, set] = defaultdict(set)
    doc_terms: List[set] = []
    for d in docs:
        terms = set(tokenize(f"{d.title} {d.snippet} {d.content[:1500]}"))
        terms = {t for t in terms if len(t) > 3}
        doc_terms.append(terms)
        for t in terms:
            df[t].add(registrable(d.domain))
    n_domains = len({registrable(d.domain) for d in docs}) or 1
    for d, terms in zip(docs, doc_terms):
        if not terms:
            d.agreement = 0.3
            continue
        # average share of other domains that mention this doc's terms
        shares = [len(df[t]) / n_domains for t in terms]
        shares.sort(reverse=True)
        top = shares[: max(8, len(shares) // 10)]
        d.agreement = min(1.0, sum(top) / len(top)) if top else 0.3
        # multi-provider consensus is a direct corroboration signal too
        d.agreement = min(1.0, d.agreement * 0.7 + 0.3 * min(1.0, len(set(d.providers)) / 4))


def _normalize(docs: Sequence[Document], field: str) -> None:
    vals = [getattr(d, field) for d in docs]
    if not vals:
        return
    lo, hi = min(vals), max(vals)
    span = hi - lo
    for d in docs:
        setattr(d, field, (getattr(d, field) - lo) / span if span > 1e-12 else
                (1.0 if hi > 0 else 0.0))


# --------------------------------------------------------------------------
# 5-6. final score + diversity
# --------------------------------------------------------------------------
def final_score(docs: Sequence[Document]) -> None:
    s = settings
    for d in docs:
        score = (
            s.w_rrf * d.rrf
            + s.w_bm25 * d.bm25
            + s.w_authority * d.authority
            + s.w_freshness * d.freshness
            + s.w_agreement * d.agreement
            + s.w_llm * d.llm_score
        )
        if d.fetched and d.word_count > 250:
            score += 0.18   # we actually read it: more trustworthy evidence
        if d.word_count and d.word_count < 60:
            score -= 0.12   # thin page
        # File results are a direct answer to a download query, not prose.
        # They can never be "read", so exempt them and reward swarm health.
        if (d.extra or {}).get("kind") in ("torrent", "repack"):
            score += 0.55
            try:
                score += min(0.45, int((d.extra or {}).get("seeders") or 0) / 1500.0)
            except (TypeError, ValueError):
                pass
        d.score = score


def diversify(
    docs: List[Document], limit: int, per_domain: int = 2, lam: float = 0.86
) -> List[Document]:
    """
    Greedy selection balancing score against redundancy: caps per-domain
    hits and penalises near-duplicate titles so the source list stays broad.
    """
    ranked = sorted(docs, key=lambda d: d.score, reverse=True)
    picked: List[Document] = []
    dom_count: Dict[str, int] = defaultdict(int)
    seen_sets: List[set] = []
    for d in ranked:
        if len(picked) >= limit:
            break
        reg = registrable(d.domain)
        if dom_count[reg] >= per_domain:
            continue
        terms = set(tokenize(d.title)[:14])
        if terms:
            dup = any(
                len(terms & prev) / max(1, len(terms | prev)) > 0.82
                for prev in seen_sets
            )
            if dup:
                continue
        picked.append(d)
        dom_count[reg] += 1
        seen_sets.append(terms)
    if len(picked) < limit:  # backfill if diversity was too strict
        for d in ranked:
            if len(picked) >= limit:
                break
            if d not in picked:
                picked.append(d)
    return picked[:limit]


def rank_documents(
    docs: List[Document],
    queries: Sequence[str],
    limit: int,
    per_domain: int = 2,
) -> List[Document]:
    if not docs:
        return []
    compute_rrf(docs, k=settings.rrf_k)
    compute_bm25(docs, queries)
    compute_authority(docs)
    compute_freshness(docs)
    compute_agreement(docs)
    final_score(docs)
    return diversify(docs, limit, per_domain=per_domain)


# --------------------------------------------------------------------------
# evidence chunk ranking
# --------------------------------------------------------------------------
def rank_chunks(chunks: List[Chunk], queries: Sequence[str], limit: int) -> List[Chunk]:
    if not chunks:
        return []
    corpus = [tokenize(c.text) for c in chunks]
    try:
        bm = BM25Okapi(corpus)
    except Exception:
        return chunks[:limit]
    agg = [0.0] * len(chunks)
    for qi, q in enumerate(queries):
        qt = tokenize(q)
        if not qt:
            continue
        w = 1.0 if qi == 0 else 0.6
        for i, s in enumerate(bm.get_scores(qt)):
            agg[i] += float(s) * w
    mx = max(agg) or 1.0
    for c, s in zip(chunks, agg):
        # blend passage relevance with its parent document's quality
        c.score = 0.72 * (s / mx) + 0.28 * min(1.0, c.doc.score / 3.0)
    ranked = sorted(chunks, key=lambda c: c.score, reverse=True)

    # keep at most 3 passages per document so one page can't flood context
    out: List[Chunk] = []
    per_doc: Dict[str, int] = defaultdict(int)
    for c in ranked:
        if per_doc[c.doc.url] >= 3:
            continue
        out.append(c)
        per_doc[c.doc.url] += 1
        if len(out) >= limit:
            break
    return out