"""Multi-lane knowledge retrieval (hybrid search, FTS fallback, topic boost, MMR fusion)."""

from __future__ import annotations

import re
from typing import Any, Literal

from core.config import get_config
from tools._logging import _debug_log

RetrievalMode = Literal["ask", "tool", "inject", "connections", "theme_detail", "unified"]

_MODE_MAX_HITS: dict[str, int] = {
    "ask": 16,
    "tool": 12,
    "inject": 8,
    "connections": 12,
    "theme_detail": 48,
    "unified": 12,
}

_MODE_SNIPPET_CHARS: dict[str, int] = {
    "ask": 1000,
    "tool": 800,
    "inject": 600,
    "connections": 800,
    "theme_detail": 600,
    "unified": 800,
}

_DEFAULT_LANE_WEIGHTS = {
    "semantic": 1.0,
    "fts_or": 0.88,
    "topic": 0.92,
    "text_fallback": 0.42,
}

_DEFINITION_RE = re.compile(
    r"(?is)\bwhat\s+(?:is|are)\s+(?:the\s+)?(.+?)(?:\?|$|\s+and\s+|\s+versus\s+|\s+vs\s+)",
)

_STALE_INDEX_STATES = frozenset(
    {"stale_chunker", "stale_identity", "vectors_missing", "reindexing"}
)


def get_retrieval_config() -> dict[str, Any]:
    """[knowledge.retrieval] with safe defaults."""
    knowledge = get_config().get("knowledge") or {}
    raw = knowledge.get("retrieval") if isinstance(knowledge.get("retrieval"), dict) else {}
    sec = dict(raw) if isinstance(raw, dict) else {}
    sec.setdefault("max_sub_queries", 3)
    sec.setdefault("ask_max_hits", 16)
    sec.setdefault("tool_max_hits", 12)
    sec.setdefault("inject_max_hits", 8)
    sec.setdefault("mmr_per_source_cap", 3)
    sec.setdefault("rare_token_min_len", 8)
    sec.setdefault("min_fused_score", 0.05)
    sec.setdefault("ask_min_fused_score", 0.10)
    sec.setdefault("inject_min_fused_score", 0.10)
    sec.setdefault("lane_weight_semantic", _DEFAULT_LANE_WEIGHTS["semantic"])
    sec.setdefault("lane_weight_fts_or", _DEFAULT_LANE_WEIGHTS["fts_or"])
    sec.setdefault("lane_weight_topic", _DEFAULT_LANE_WEIGHTS["topic"])
    sec.setdefault("lane_weight_text_fallback", _DEFAULT_LANE_WEIGHTS["text_fallback"])
    sec.setdefault("allow_text_fallback", True)
    return sec


def _cfg_float(cfg: dict[str, Any], key: str, default: float) -> float:
    raw = cfg.get(key)
    if raw is None:
        return default
    try:
        return float(raw)
    except (TypeError, ValueError):
        return default


def _cfg_int(cfg: dict[str, Any], key: str, default: int) -> int:
    raw = cfg.get(key)
    if raw is None:
        return default
    try:
        return int(raw)
    except (TypeError, ValueError):
        return default


def _cfg_bool(cfg: dict[str, Any], key: str, default: bool) -> bool:
    raw = cfg.get(key, default)
    if isinstance(raw, bool):
        return raw
    if isinstance(raw, (int, float)):
        return bool(raw)
    if isinstance(raw, str):
        low = raw.strip().lower()
        if low in ("0", "false", "no", "off"):
            return False
        if low in ("1", "true", "yes", "on"):
            return True
    return default


def ask_fused_score_floor(cfg: dict[str, Any] | None = None) -> float:
    """Fused-score floor for Library Ask (stricter than global ``min_fused_score`` by default)."""
    sec = cfg if cfg is not None else get_retrieval_config()
    ask_floor = sec.get("ask_min_fused_score")
    if ask_floor is not None:
        try:
            return max(0.0, float(ask_floor))
        except (TypeError, ValueError):
            pass
    return max(0.0, _cfg_float(sec, "min_fused_score", 0.05))


def inject_fused_score_floor(cfg: dict[str, Any] | None = None) -> float:
    """Fused-score floor for ambient knowledge injection.

    Mirrors :func:`ask_fused_score_floor`: ambient injection runs on *every*
    foreground turn without the user asking for it, so it must be at least as
    strict as Library Ask to avoid dumping weakly-related passages into the
    system prompt on off-topic conversations.
    """
    sec = cfg if cfg is not None else get_retrieval_config()
    inject_floor = sec.get("inject_min_fused_score")
    if inject_floor is not None:
        try:
            return max(0.0, float(inject_floor))
        except (TypeError, ValueError):
            pass
    return ask_fused_score_floor(sec)


def _lane_weights(cfg: dict[str, Any]) -> dict[str, float]:
    def _w(key: str, default: float) -> float:
        return max(0.0, _cfg_float(cfg, key, default))

    return {
        "semantic": _w("lane_weight_semantic", _DEFAULT_LANE_WEIGHTS["semantic"]),
        "fts_or": _w("lane_weight_fts_or", _DEFAULT_LANE_WEIGHTS["fts_or"]),
        "topic": _w("lane_weight_topic", _DEFAULT_LANE_WEIGHTS["topic"]),
        "text_fallback": _w("lane_weight_text_fallback", _DEFAULT_LANE_WEIGHTS["text_fallback"]),
    }


def _merge_pick_better(
    prev: dict[str, Any] | None, hit: dict[str, Any]
) -> dict[str, Any]:
    if prev is None:
        return hit
    prev_f = float(prev.get("fused_score") or 0)
    hit_f = float(hit.get("fused_score") or 0)
    if hit_f > prev_f:
        return hit
    if hit_f < prev_f:
        return prev
    if float(hit.get("base_score") or 0) > float(prev.get("base_score") or 0):
        return hit
    return prev


def _search_row_base_score(row: dict[str, Any]) -> float:
    """Raw relevance from a session_search_db row (cosine, RRF, or FTS rank)."""
    sem = row.get("semantic_score")
    if sem is not None:
        try:
            v = float(sem)
            if v > 0:
                return v
        except (TypeError, ValueError):
            pass
    rrf = row.get("rrf_score")
    if rrf is not None:
        try:
            v = float(rrf)
            if v > 0:
                return v
        except (TypeError, ValueError):
            pass
    legacy = row.get("score")
    if legacy is not None:
        try:
            v = float(legacy)
            if v > 0:
                return v
        except (TypeError, ValueError):
            pass
    rank = row.get("rank")
    try:
        r = abs(float(rank))
        return 1.0 / (1.0 + r)
    except (TypeError, ValueError):
        pass
    return 0.0


def _normalize_batch_relevance(base_scores: list[float]) -> list[float]:
    """Rank-order base scores within one lane batch to comparable [0, 1] relevances."""
    n = len(base_scores)
    if n == 0:
        return []
    if n == 1:
        return [1.0]
    order = sorted(range(n), key=lambda i: base_scores[i], reverse=True)
    rel = [0.0] * n
    for rank_idx, orig_i in enumerate(order):
        rel[orig_i] = 1.0 - (rank_idx / (n - 1))
    return rel


def _assign_lane_fused_scores(
    hits: list[dict[str, Any]],
    *,
    lane: str,
    weights: dict[str, float],
) -> None:
    """Apply lane_weight * rank-normalized relevance so MMR orders by true ranking."""
    lane_w = weights.get(lane, weights.get("semantic", 1.0))
    bases = [float(h.get("base_score") or 0) for h in hits]
    rels = _normalize_batch_relevance(bases)
    for hit, rel in zip(hits, rels):
        hit["fused_score"] = lane_w * rel


def _cap_for_mode(mode: RetrievalMode, requested: int, cfg: dict[str, Any]) -> int:
    key = {
        "ask": "ask_max_hits",
        "tool": "tool_max_hits",
        "inject": "inject_max_hits",
    }.get(mode)
    default = _MODE_MAX_HITS.get(mode, 12)
    if key and key in cfg:
        default = max(1, min(_cfg_int(cfg, key, default), 20))
    hard_cap = 48 if mode == "theme_detail" else 20
    try:
        req = int(requested)
    except (TypeError, ValueError):
        req = default
    return max(1, min(req, default, hard_cap))


def _plan_sub_queries(question: str, cfg: dict[str, Any]) -> list[tuple[str, str]]:
    """Return (sub_query, lane_hint) pairs, bounded by max_sub_queries."""
    from tools.knowledge_tools import _tokenizer_refs

    q = (question or "").strip()
    if not q:
        return []
    max_sq = max(1, min(_cfg_int(cfg, "max_sub_queries", 3), 3))
    planned: list[tuple[str, str]] = [("semantic", q)]

    token_re, stopwords = _tokenizer_refs()
    tokens = [
        t.lower()
        for t in token_re.findall(q.lower())
        if len(t) >= 3 and t.lower() not in stopwords
    ]
    if tokens:
        or_query = " OR ".join(dict.fromkeys(tokens))
        planned.append(("fts_or", or_query))

    m = _DEFINITION_RE.search(q)
    if m:
        term = m.group(1).strip(" .,;:\"'")
        if len(term) >= 3 and term.lower() not in stopwords:
            planned.append(("semantic", term[:200]))

    out: list[tuple[str, str]] = []
    seen_q: set[str] = set()
    for lane, sub in planned:
        key = sub.lower()
        if key in seen_q:
            continue
        seen_q.add(key)
        out.append((lane, sub))
        if len(out) >= max_sq:
            break
    return out


def _topic_terms_for_corpus(slug: str) -> set[str]:
    from tools.knowledge_tools import _load_manifest

    manifest = _load_manifest(slug) or {}
    terms: set[str] = set()
    for t in (manifest.get("topics") or []) + (manifest.get("broad_topics") or []):
        s = str(t).strip().lower()
        if s:
            terms.add(s)
    return terms


def _question_has_rare_token(question: str, corpus_terms: set[str], cfg: dict[str, Any]) -> bool:
    from tools.knowledge_tools import _tokenizer_refs

    min_len = max(4, _cfg_int(cfg, "rare_token_min_len", 8))
    token_re, stopwords = _tokenizer_refs()
    for t in token_re.findall(question.lower()):
        if t.lower() in stopwords or len(t) < 3:
            continue
        if len(t) >= min_len and t.lower() not in corpus_terms:
            return True
    return False


def _topic_lane_queries(question: str, slug: str) -> list[str]:
    q_low = question.lower()
    manifest_terms = _topic_terms_for_corpus(slug)
    hits: list[str] = []
    for term in sorted(manifest_terms, key=len, reverse=True):
        if len(term) >= 4 and term in q_low:
            hits.append(term)
    return hits[:2]


def _lane_search(
    sub_query: str,
    *,
    slug: str,
    label: str,
    thread: str,
    lane: str,
    cap: int,
    semantic: bool,
    force_or: bool,
    weights: dict[str, float],
) -> list[dict[str, Any]]:
    from tools.knowledge_tools import _hit_from_search_row

    try:
        import session_search_db
    except Exception as exc:
        _debug_log(f"[knowledge] retrieval: session_search_db unavailable: {exc}", "WARN")
        return []

    fts_modes: list[str] = []
    if force_or or lane == "fts_or":
        fts_modes = ["or"]
    else:
        fts_modes = ["and_content", "or"]

    rows: list[dict[str, Any]] = []
    for fts_mode in fts_modes:
        rows = session_search_db.search(
            sub_query,
            exclude_thread_id="",
            max_results=cap,
            semantic=semantic and lane != "fts_or",
            hybrid=True,
            include_thread_id=thread,
            fts_mode=fts_mode,
        )
        if rows:
            break

    out: list[dict[str, Any]] = []
    for row in rows:
        hit = _hit_from_search_row(row, slug, label)
        if not hit:
            continue
        hit["retrieval_lane"] = lane
        hit["base_score"] = _search_row_base_score(row)
        hit["search_source"] = str(row.get("_source") or "")
        out.append(hit)
    lane_key = "fts_or" if lane == "fts_or" else lane
    _assign_lane_fused_scores(out, lane=lane_key, weights=weights)
    return out


def _corpus_semantic_index_ready(slug: str) -> bool:
    """True when per-corpus semantic vectors are searchable (ready state only)."""
    return _corpus_semantic_index_state(slug) == "ready"


def _global_embedding_reindex_state() -> str:
    try:
        from utils.embedding_model import get_embedding_reindex_status

        return str(get_embedding_reindex_status().get("state") or "idle")
    except Exception:
        return "idle"


def _corpus_semantic_index_state(slug: str) -> str:
    try:
        from utils.embedding_chunking import corpus_semantic_index_state

        return corpus_semantic_index_state(
            f"knowledge:{slug}",
            embedding_reindex_state=_global_embedding_reindex_state(),
        )
    except Exception as exc:
        _debug_log(
            f"[knowledge] retrieval: semantic index state check failed for {slug}: {exc}",
            "WARN",
        )
        return "unknown"


def _corpus_allows_text_fallback(slug: str) -> bool:
    """Fail closed on stale/missing vectors; allow FTS-only and ready corpora."""
    return _corpus_semantic_index_state(slug) in ("ready", "fts_only")


def _should_use_text_fallback(
    *,
    slug: str,
    semantic: bool,
    mode: RetrievalMode,
    planned_lane: str,
    batch_empty: bool,
    cfg: dict[str, Any],
) -> bool:
    # Primary semantic sub-query only (includes rare-token force_or reroute to fts_or).
    # Caller `semantic` is ignored; per-corpus index state gates raw-file grep.
    _ = semantic
    if mode == "inject" or planned_lane != "semantic" or not batch_empty:
        return False
    if not _cfg_bool(cfg, "allow_text_fallback", True):
        return False
    return _corpus_allows_text_fallback(slug)


def _text_fallback_hits(
    sub_query: str,
    *,
    slug: str,
    label: str,
    cap: int,
    weights: dict[str, float],
) -> list[dict[str, Any]]:
    from tools.knowledge_tools import (
        _load_manifest,
        _read_text_file,
        _search_in_text,
        _source_page_path,
        _synthesis_path,
    )

    import os

    def _fallback_looks_secret(text: str) -> bool:
        lower = (text or "").lower()
        if any(
            marker in lower
            for marker in ("api_key", "apikey", "secret", "password", "authorization:", "bearer ")
        ):
            return True
        return "sk-" in text and len(text) > 24

    try:
        from chat_storage import _looks_secret
    except Exception:
        _looks_secret = _fallback_looks_secret  # type: ignore[assignment]

    out: list[dict[str, Any]] = []

    def _append_fallback_hit(entry: dict[str, Any]) -> None:
        entry["base_score"] = max(0.01, 1.0 - (len(out) * 0.05))
        out.append(entry)

    synthesis = _read_text_file(_synthesis_path(slug)) if os.path.isfile(_synthesis_path(slug)) else ""
    for text_hit in _search_in_text(sub_query, synthesis):
        if _looks_secret(text_hit):
            continue
        _append_fallback_hit(
            {
                "id": f"{slug}#synthesis#{hash(text_hit) & 0xFFFFFFFF}",
                "corpus_slug": slug,
                "corpus_label": label,
                "source_slug": "",
                "source_name": "synthesis",
                "snippet": text_hit,
                "thread_id": f"knowledge:{slug}",
                "score": None,
                "retrieval_lane": "text_fallback",
            }
        )
        if len(out) >= cap:
            _assign_lane_fused_scores(out, lane="text_fallback", weights=weights)
            return out
    manifest = _load_manifest(slug) or {}
    for src in manifest.get("sources") or []:
        if not isinstance(src, dict):
            continue
        source_slug = str(src.get("slug") or "")
        source_name = str(src.get("filename") or source_slug)
        path = _source_page_path(slug, source_slug)
        if not os.path.isfile(path):
            continue
        content = _read_text_file(path)
        for text_hit in _search_in_text(sub_query, content):
            if _looks_secret(text_hit):
                continue
            _append_fallback_hit(
                {
                    "id": f"{slug}#{source_slug}#{hash(text_hit) & 0xFFFFFFFF}",
                    "corpus_slug": slug,
                    "corpus_label": label,
                    "source_slug": source_slug,
                    "source_name": source_name,
                    "snippet": text_hit,
                    "thread_id": f"knowledge:{slug}",
                    "score": None,
                    "retrieval_lane": "text_fallback",
                }
            )
            if len(out) >= cap:
                _assign_lane_fused_scores(out, lane="text_fallback", weights=weights)
                return out
    _assign_lane_fused_scores(out, lane="text_fallback", weights=weights)
    return out


def _trim_snippet(hit: dict[str, Any], max_chars: int) -> None:
    snip = str(hit.get("snippet") or "")
    if len(snip) > max_chars:
        hit["snippet"] = snip[: max_chars - 1] + "…"


def _apply_mmr(
    hits: list[dict[str, Any]],
    *,
    cap: int,
    per_source: int,
) -> list[dict[str, Any]]:
    """Greedy MMR by fused_score with per-source_slug cap."""
    pool = sorted(
        hits,
        key=lambda h: (
            float(h.get("fused_score") or 0),
            float(h.get("base_score") or 0),
        ),
        reverse=True,
    )
    chosen: list[dict[str, Any]] = []
    per_src: dict[str, int] = {}
    for hit in pool:
        src = "|".join(
            [
                str(hit.get("corpus_slug") or ""),
                str(hit.get("source_slug") or hit.get("source_name") or "synthesis"),
            ]
        )
        if per_src.get(src, 0) >= per_source:
            continue
        chosen.append(hit)
        per_src[src] = per_src.get(src, 0) + 1
        if len(chosen) >= cap:
            break
    return chosen


def retrieve_knowledge_hits(
    query: str,
    corpus: str | None = None,
    max_results: int = 8,
    *,
    mode: RetrievalMode = "tool",
    semantic: bool = True,
) -> list[dict[str, Any]]:
    """Fused multi-lane retrieval over digested knowledge corpora."""
    from tools.knowledge_tools import _filter_corpora_by_name, list_knowledge_corpora

    q = (query or "").strip()
    if not q:
        return []

    cfg = get_retrieval_config()
    cap = _cap_for_mode(mode, max_results, cfg)
    snippet_cap = _MODE_SNIPPET_CHARS.get(mode, 800)
    per_source = max(1, min(_cfg_int(cfg, "mmr_per_source_cap", 3), 8))
    if mode == "theme_detail":
        per_source = 1

    corpora = _filter_corpora_by_name(list_knowledge_corpora(), corpus)
    if not corpora:
        return []

    sub_queries = _plan_sub_queries(q, cfg)
    weights = _lane_weights(cfg)
    merged: dict[str, dict[str, Any]] = {}

    for c in corpora:
        slug = str(c["slug"])
        label = str(c["label"])
        thread = f"knowledge:{slug}"
        corpus_terms = _topic_terms_for_corpus(slug)
        force_or = _question_has_rare_token(q, corpus_terms, cfg)

        for lane, sub in sub_queries:
            lane_key = "fts_or" if lane == "fts_or" or (lane == "semantic" and force_or) else lane
            if lane_key == "fts_or":
                batch = _lane_search(
                    sub,
                    slug=slug,
                    label=label,
                    thread=thread,
                    lane="fts_or",
                    cap=cap,
                    semantic=False,
                    force_or=True,
                    weights=weights,
                )
            else:
                batch = _lane_search(
                    sub,
                    slug=slug,
                    label=label,
                    thread=thread,
                    lane="semantic",
                    cap=cap,
                    semantic=semantic,
                    force_or=False,
                    weights=weights,
                )
            if _should_use_text_fallback(
                slug=slug,
                semantic=semantic,
                mode=mode,
                planned_lane=lane,
                batch_empty=not batch,
                cfg=cfg,
            ):
                batch = _text_fallback_hits(
                    sub, slug=slug, label=label, cap=cap, weights=weights
                )
            for hit in batch:
                hid = str(hit.get("id") or "")
                if not hid:
                    continue
                merged[hid] = _merge_pick_better(merged.get(hid), hit)

        for topic_q in _topic_lane_queries(q, slug):
            topic_batch = _lane_search(
                topic_q,
                slug=slug,
                label=label,
                thread=thread,
                lane="topic",
                cap=min(4, cap),
                semantic=semantic,
                force_or=False,
                weights=weights,
            )
            for hit in topic_batch:
                hid = str(hit.get("id") or "")
                if not hid:
                    continue
                merged[hid] = _merge_pick_better(merged.get(hid), hit)

        if not any(str(h.get("corpus_slug") or "") == slug for h in merged.values()):
            state = _corpus_semantic_index_state(slug)
            if state in _STALE_INDEX_STATES:
                try:
                    from utils.embedding_model import maybe_schedule_embedding_reindex

                    maybe_schedule_embedding_reindex(f"knowledge:{slug}")
                except Exception as exc:
                    _debug_log(
                        f"[knowledge] retrieval: reindex schedule failed for {slug}: {exc}",
                        "WARN",
                    )

    if not merged:
        _debug_log(f"[knowledge] retrieval: no hits for query ({mode})", "INFO")
        return []

    fused = _apply_mmr(list(merged.values()), cap=cap, per_source=per_source)
    if len(fused) < 3:
        _debug_log(
            f"[knowledge] retrieval: only {len(fused)} hit(s) after fusion ({mode})",
            "INFO",
        )
    for hit in fused:
        _trim_snippet(hit, snippet_cap)
    return fused


def collect_retrieval_metrics(hits: list[dict[str, Any]]) -> dict[str, Any]:
    """Hit@1 diagnostics for eval/smoke (not golden-snapshotted)."""
    ranked = sorted(
        hits,
        key=lambda h: (
            float(h.get("fused_score") or 0),
            float(h.get("base_score") or 0),
        ),
        reverse=True,
    )
    lanes: dict[str, int] = {}
    for h in ranked:
        lane = str(h.get("retrieval_lane") or "unknown")
        lanes[lane] = lanes.get(lane, 0) + 1
    top = ranked[0] if ranked else {}
    return {
        "hit_count": len(hits),
        "lanes": lanes,
        "top_source_slug": top.get("source_slug"),
        "top_fused_score": top.get("fused_score"),
        "fused_scores": [h.get("fused_score") for h in ranked[:5]],
    }


def retrieval_lanes_used(hits: list[dict[str, Any]]) -> list[str]:
    lanes: list[str] = []
    seen: set[str] = set()
    for h in hits:
        lane = str(h.get("retrieval_lane") or "")
        if lane and lane not in seen:
            seen.add(lane)
            lanes.append(lane)
    return lanes
