"""Multi-path merge for long-term memory search (lexical + TF-IDF chunks + recency).

Used by ``search_memory`` in ``tools/memory_tools.py``. Config keys live under
``[memory]`` in mugi_config.toml (see ``get_memory_config`` setdefaults).
"""
from __future__ import annotations

import math
import os
import re
from dataclasses import dataclass
from datetime import date
from typing import Any

# Defaults merged into [memory] via core.config.get_memory_config
_FUSION_DEFAULTS: dict[str, Any] = {
    "memory_search_max_bytes_per_file": 200_000,
    "memory_search_max_total_bytes": 1_500_000,
    "memory_search_max_shards_scan": 200,
    "memory_search_max_semantic_chunks": 600,
    "memory_search_fusion_weight_lexical": 1.0,
    "memory_search_fusion_weight_semantic": 0.85,
    "memory_search_fusion_weight_recency": 0.35,
    "memory_search_per_source_char_cap": 1_200,
    "memory_search_semantic_min_score": 0.05,
}


def apply_fusion_config_defaults(mem: dict[str, Any]) -> dict[str, Any]:
    """Mutate memory section dict with fusion keys; return same dict."""
    for k, v in _FUSION_DEFAULTS.items():
        mem.setdefault(k, v)
    # Coerce numeric types from TOML
    try:
        mem["memory_search_max_bytes_per_file"] = int(mem["memory_search_max_bytes_per_file"])
        mem["memory_search_max_total_bytes"] = int(mem["memory_search_max_total_bytes"])
        mem["memory_search_max_shards_scan"] = int(mem["memory_search_max_shards_scan"])
        mem["memory_search_max_semantic_chunks"] = int(mem["memory_search_max_semantic_chunks"])
        mem["memory_search_fusion_weight_lexical"] = float(mem["memory_search_fusion_weight_lexical"])
        mem["memory_search_fusion_weight_semantic"] = float(mem["memory_search_fusion_weight_semantic"])
        mem["memory_search_fusion_weight_recency"] = float(mem["memory_search_fusion_weight_recency"])
        mem["memory_search_per_source_char_cap"] = int(mem["memory_search_per_source_char_cap"])
        mem["memory_search_semantic_min_score"] = float(mem["memory_search_semantic_min_score"])
    except (TypeError, ValueError):
        pass
    return mem


def apply_memory_retrieval_profile_to_fusion_weights(mem: dict[str, Any], profile: str | None) -> None:
    """Bias fusion weights from app pref ``memory_retrieval_profile`` (conservative | balanced)."""
    p = (profile or "").strip().lower()
    if p != "conservative":
        return
    try:
        w_lex = float(mem.get("memory_search_fusion_weight_lexical", 1.0))
        w_sem = float(mem.get("memory_search_fusion_weight_semantic", 0.85))
    except (TypeError, ValueError):
        return
    mem["memory_search_fusion_weight_lexical"] = max(w_lex, 1.18)
    mem["memory_search_fusion_weight_semantic"] = min(w_sem, 0.70)


_SUMMARY_NAME_RE = re.compile(r"^(?:memory/summaries/)?(\d{4})-(\d{2})-(\d{2})\.md$")
_ARCHIVES_RE = re.compile(r"memory/archives/")


def recency_for_source(source_label: str, *, today: date | None = None) -> float:
    """0..1 recency bump: fresh daily summaries highest; hub/shards middling; archives low."""
    d = today or date.today()
    s = (source_label or "").replace("\\", "/")
    m = _SUMMARY_NAME_RE.search(s)
    if m:
        try:
            file_date = date(int(m.group(1)), int(m.group(2)), int(m.group(3)))
            days = max(0, (d - file_date).days)
            return max(0.0, 1.0 - min(days, 90) / 90.0)
        except ValueError:
            pass
    if "MEMORY.md" == s.split("/")[-1] or s.endswith("MEMORY.md"):
        return 0.55
    if s.startswith("memory/shards/"):
        return 0.5
    if _ARCHIVES_RE.search(s):
        return 0.2
    return 0.45


def _tokenize(text: str) -> list[str]:
    return [t for t in re.split(r"[^a-z0-9_]+", (text or "").lower()) if t]


def _token_overlap_score(query: str, text: str) -> float:
    q = set(_tokenize(query))
    t = set(_tokenize(text))
    if not q or not t:
        return 0.0
    inter = len(q & t)
    denom = math.sqrt(len(q)) * math.sqrt(len(t))
    return inter / denom if denom else 0.0


def _jaccard_tokens(a: str, b: str) -> float:
    sa, sb = set(_tokenize(a)), set(_tokenize(b))
    if not sa and not sb:
        return 1.0
    if not sa or not sb:
        return 0.0
    return len(sa & sb) / len(sa | sb)


def _tfidf_cosine_scores(query: str, docs: list[tuple[str, str]]) -> list[float]:
    """Parallel to memory_tools._semantic_hits: one score per doc (same order as docs)."""
    q_tokens = _tokenize(query)
    if not q_tokens:
        return [0.0] * len(docs)
    doc_tokens: list[tuple[list[str], str]] = []
    for src, txt in docs:
        doc_tokens.append((_tokenize(txt), txt))
    if not doc_tokens:
        return []
    n_docs = len(doc_tokens)
    df: dict[str, int] = {}
    for toks, _ in doc_tokens:
        for tok in set(toks):
            df[tok] = df.get(tok, 0) + 1
    idf = {tok: math.log((1 + n_docs) / (1 + cnt)) + 1.0 for tok, cnt in df.items()}

    def vec(tokens: list[str]) -> dict[str, float]:
        tf: dict[str, float] = {}
        for tok in tokens:
            tf[tok] = tf.get(tok, 0.0) + 1.0
        for tok in list(tf.keys()):
            tf[tok] = tf[tok] * idf.get(tok, 1.0)
        return tf

    qv = vec(q_tokens)
    qnorm = math.sqrt(sum(v * v for v in qv.values())) or 1.0
    scores: list[float] = []
    for toks, _ in doc_tokens:
        dv = vec(toks)
        dnorm = math.sqrt(sum(v * v for v in dv.values())) or 1.0
        dot = 0.0
        for tok, w in qv.items():
            dot += w * dv.get(tok, 0.0)
        scores.append(dot / (qnorm * dnorm))
    return scores


@dataclass
class _Cand:
    source: str
    text: str
    score: float


def fuse_memory_search(
    query: str,
    *,
    lexical_hits: list[tuple[str, str, float]],
    semantic_docs: list[tuple[str, str, float]],
    cfg: dict[str, Any],
    max_results: int,
) -> tuple[list[str], dict[str, Any]]:
    """Merge lexical hits and semantic chunks into ranked, deduplicated excerpts.

    ``lexical_hits`` / ``semantic_docs`` entries are ``(source_label, text, recency)``.

    Returns ``(formatted_blocks, stats)`` where each block is
    ``"[source]\\nsnippet"`` and stats includes merge telemetry.
    """
    w_lex = float(cfg.get("memory_search_fusion_weight_lexical", 1.0))
    w_sem = float(cfg.get("memory_search_fusion_weight_semantic", 0.85))
    w_rec = float(cfg.get("memory_search_fusion_weight_recency", 0.35))
    sem_min = float(cfg.get("memory_search_semantic_min_score", 0.05))
    per_source_cap = int(cfg.get("memory_search_per_source_char_cap", 1200))

    stats: dict[str, Any] = {
        "lexical_in": len(lexical_hits),
        "semantic_chunks_in": len(semantic_docs),
        "candidates_after_merge": 0,
        "dedup_dropped": 0,
        "per_source_trimmed": 0,
    }

    if not lexical_hits and not semantic_docs:
        return [], stats

    # Semantic scores for all chunks
    sem_pairs = [(s, t) for s, t, _ in semantic_docs]
    sem_scores = _tfidf_cosine_scores(query, sem_pairs) if sem_pairs else []

    def best_sem_for_snippet(source: str, snippet: str) -> float:
        best = 0.0
        for i, (src, chunk, _) in enumerate(semantic_docs):
            if src != source:
                continue
            if snippet.strip() and (snippet in chunk or chunk in snippet):
                best = max(best, sem_scores[i])
            else:
                jac = _jaccard_tokens(snippet, chunk)
                if jac > 0.35:
                    best = max(best, sem_scores[i] * jac)
        if best == 0.0:
            o = _token_overlap_score(query, snippet)
            for i, (src, _, _) in enumerate(semantic_docs):
                if src == source:
                    best = max(best, sem_scores[i] * max(0.25, o))
                    break
            best = max(best, o * 0.9)
        return best

    candidates: list[_Cand] = []

    seen_lexical: set[tuple[str, str]] = set()
    for src, snippet, rec in lexical_hits:
        key = (src, snippet.strip()[:240])
        if key in seen_lexical:
            continue
        seen_lexical.add(key)
        bs = best_sem_for_snippet(src, snippet)
        ov = _token_overlap_score(query, snippet)
        sem_part = max(bs, ov * 0.95)
        combined = w_lex * 1.0 + w_sem * sem_part + w_rec * rec
        candidates.append(_Cand(src, snippet, combined))

    # Semantic-only: chunks above threshold not redundant with an existing candidate
    covered_chunks: set[int] = set()
    for i, c in enumerate(candidates):
        for j, (src, chunk, rec) in enumerate(semantic_docs):
            if src != c.source:
                continue
            if c.text in chunk or chunk in c.text or _jaccard_tokens(c.text, chunk) > 0.55:
                covered_chunks.add(j)

    for j, (src, chunk, rec) in enumerate(semantic_docs):
        if j in covered_chunks:
            continue
        sc = sem_scores[j] if j < len(sem_scores) else 0.0
        if sc < sem_min:
            continue
        combined = w_sem * sc + w_rec * rec
        candidates.append(_Cand(src, chunk, combined))

    candidates.sort(key=lambda c: c.score, reverse=True)
    stats["candidates_after_merge"] = len(candidates)

    # Dedupe greedy (high score first)
    kept: list[_Cand] = []
    for c in candidates:
        dup = False
        for k in kept:
            if c.source == k.source and _jaccard_tokens(c.text, k.text) > 0.82:
                dup = True
                stats["dedup_dropped"] += 1
                break
            if _jaccard_tokens(c.text, k.text) > 0.88:
                dup = True
                stats["dedup_dropped"] += 1
                break
        if not dup:
            kept.append(c)

    # Per-source char budget + max_results
    source_chars: dict[str, int] = {}
    out_lines: list[str] = []
    for c in kept:
        if len(out_lines) >= max_results:
            break
        used = source_chars.get(c.source, 0)
        if used >= per_source_cap:
            stats["per_source_trimmed"] += 1
            continue
        body = c.text.strip()
        room = per_source_cap - used
        if len(body) > room:
            body = body[: max(80, room - 3)] + "..."
            stats["per_source_trimmed"] += 1
        source_chars[c.source] = used + len(body)
        out_lines.append(f"[{c.source}]\n{body}")

    return out_lines, stats


def read_memory_file_capped(abs_path: str, max_bytes: int) -> tuple[str, bool]:
    """Read up to ``max_bytes`` of **file** UTF-8 (raw bytes), then decode.

    Uses a byte cap so this matches ``memory_search_max_*_bytes`` config and the
    caller's ``len(text.encode('utf-8'))`` accounting. Trailing bytes that would
    form an incomplete code point are dropped before decode.
    """
    if max_bytes <= 0:
        return "", False
    if not abs_path or not os.path.isfile(abs_path):
        return "", False
    try:
        with open(abs_path, "rb") as f:
            raw = f.read(max_bytes + 1)
        truncated = len(raw) > max_bytes
        if truncated:
            raw = raw[:max_bytes]
        # Drop an incomplete UTF-8 code point at the cut (suffix is at most 3 bytes).
        for _ in range(4):
            if not raw:
                break
            try:
                raw.decode("utf-8")
                break
            except UnicodeDecodeError:
                raw = raw[:-1]
        return raw.decode("utf-8", errors="replace"), truncated
    except OSError:
        return "", False
