"""Lexical + vector search helpers for ``knowtation/search`` (§KD-5a.A(a))."""

from __future__ import annotations

import json
import logging
import os
import pathlib
import re
import urllib.error
import urllib.parse
import urllib.request
from typing import Any, TypedDict

from muse.plugins.knowtation.hooks import DEFAULT_PORT, ENV_PORT, _resolve_port
from muse.plugins.knowtation.link_index import _MAX_NOTE_BYTES, build_link_index
from muse.plugins.knowtation.mcp_helpers import _note_body_and_frontmatter
from muse.plugins.knowtation.parser import parse_frontmatter
from muse.plugins.knowtation.plugin import _ALWAYS_IGNORE_DIRS
from muse.plugins.knowtation.stats import collect_vault_stats

logger = logging.getLogger(__name__)

W_LEX: float = 0.4
W_VEC: float = 0.6
VECTOR_TIMEOUT_SECONDS: float = 2.0
_MAX_HUB_RESPONSE_BYTES: int = 256 * 1024

_MARKDOWN_EXTS: frozenset[str] = frozenset({".md", ".markdown", ".mdx"})
_IGNORE_DIRS: frozenset[str] = _ALWAYS_IGNORE_DIRS | frozenset({"templates", "meta"})


class LexicalHit(TypedDict, total=False):
    path: str
    score: float
    section_address: str
    snippet: str


def _note_ext(path: str) -> bool:
    lower = path.lower()
    return any(lower.endswith(ext) for ext in _MARKDOWN_EXTS)


def _query_tokens(query: str) -> list[str]:
    return [t for t in re.split(r"\s+", query.strip().lower()) if t]


def _score_text(text: str, tokens: list[str]) -> float:
    if not tokens:
        return 0.0
    hay = text.lower()
    hits = sum(1 for t in tokens if t in hay)
    return hits / len(tokens)


def lexical_search(
    root: pathlib.Path,
    query: str,
    *,
    top_k: int = 10,
) -> list[LexicalHit]:
    """Grep vault notes for *query* — title, headings, frontmatter, body."""
    tokens = _query_tokens(query)
    if not tokens:
        return []

    results: list[LexicalHit] = []
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = [d for d in dirnames if d not in _IGNORE_DIRS and not d.startswith(".")]
        for name in filenames:
            if not _note_ext(name):
                continue
            full = pathlib.Path(dirpath) / name
            try:
                rel = full.relative_to(root).as_posix()
            except ValueError:
                continue
            try:
                size = full.stat().st_size
            except OSError:
                continue
            if size > _MAX_NOTE_BYTES:
                continue
            try:
                raw = full.read_bytes()
            except OSError:
                continue
            try:
                text = raw.decode("utf-8")
            except UnicodeDecodeError:
                continue

            fm_parsed = parse_frontmatter(raw)
            title = pathlib.Path(rel).stem
            project = None
            source_type = None
            tags = None
            if fm_parsed is not None:
                title = str(fm_parsed.title or title)
                project = fm_parsed.project
                tags = fm_parsed.tags
                source_type = fm_parsed.effective_source
            if text.startswith("---"):
                parts = text.split("---", 2)
                body = parts[2] if len(parts) >= 3 else text
            else:
                body = text
            fm_text = text.split("---", 2)[1] if text.startswith("---") and len(text.split("---", 2)) >= 2 else ""
            heading_hits = re.findall(r"^#{1,6}\s+(.+)$", body, re.MULTILINE)
            combined = " ".join([title, fm_text, " ".join(heading_hits), body[:8000]])
            score = _score_text(combined, tokens)
            if score <= 0:
                continue
            snippet_src = body[:300] if body else title
            results.append(
                LexicalHit(
                    path=rel,
                    score=min(1.0, score),
                    snippet=snippet_src.replace("\n", " ")[:240],
                )
            )

    results.sort(key=lambda h: (-float(h.get("score", 0)), h.get("path", "")))
    return results[:top_k]


def vector_search(
    query: str,
    *,
    top_k: int = 10,
    port: int | None = None,
) -> tuple[list[dict[str, Any]], str]:
    """Query Knowtation hub vector index via GET (§KD-5a.A(a)).

    Returns:
        ``(hits, status)`` where *status* is ``ok`` or ``unavailable``.
    """
    resolved_port = port
    if resolved_port is None:
        p, err = _resolve_port()
        if err or p is None:
            return [], "unavailable"
        resolved_port = p

    params = urllib.parse.urlencode({"q": query, "k": str(top_k)})
    url = f"http://localhost:{resolved_port}/api/v1/search?{params}"
    request = urllib.request.Request(
        url=url,
        method="GET",
        headers={"User-Agent": "muse-mcp-search/1.0", "Accept": "application/json"},
    )
    try:
        with urllib.request.urlopen(request, timeout=VECTOR_TIMEOUT_SECONDS) as resp:  # noqa: S310
            raw = resp.read(_MAX_HUB_RESPONSE_BYTES + 1)
            if len(raw) > _MAX_HUB_RESPONSE_BYTES:
                return [], "unavailable"
            data = json.loads(raw.decode("utf-8"))
    except (urllib.error.URLError, TimeoutError, json.JSONDecodeError, OSError) as exc:
        logger.debug("vector search unavailable: %s", type(exc).__name__)
        return [], "unavailable"

    hits: list[dict[str, Any]] = []
    items = data.get("results") or data.get("hits") or data
    if isinstance(items, list):
        for item in items:
            if not isinstance(item, dict):
                continue
            path = item.get("path") or item.get("note_path") or item.get("id")
            if not isinstance(path, str) or not path:
                continue
            score_raw = item.get("score") or item.get("similarity") or 0.0
            try:
                score = float(score_raw)
            except (TypeError, ValueError):
                score = 0.0
            snippet = str(item.get("snippet") or item.get("excerpt") or "")[:240]
            hits.append({"path": path.replace("\\", "/"), "score": score, "snippet": snippet})
    return hits, "ok"


def _head_mtime(root: pathlib.Path, path: str) -> float:
    try:
        return (root / path).stat().st_mtime
    except OSError:
        return 0.0


def hybrid_search(
    root: pathlib.Path,
    query: str,
    *,
    top_k: int = 10,
    mode: str = "hybrid",
    project: str | None = None,
) -> dict[str, Any]:
    """Fan-out lexical + vector backends and merge (§KD-5a.A(a))."""
    stats = collect_vault_stats(root)
    index_freshness = stats.get("index_freshness")

    lexical_status = "ok"
    vector_status = "ok"
    degraded = False
    deduped = 0

    lex_hits: list[LexicalHit] = []
    vec_hits: list[dict[str, Any]] = []

    if mode in ("hybrid", "lexical"):
        try:
            lex_hits = lexical_search(root, query, top_k=top_k * 2)
        except Exception:
            lexical_status = "error"
            lex_hits = []

    if mode in ("hybrid", "semantic"):
        vec_hits, vs = vector_search(query, top_k=top_k * 2)
        vector_status = vs
        if vs != "ok":
            degraded = True

    if mode == "semantic" and vector_status != "ok":
        degraded = True
    if mode == "hybrid" and vector_status != "ok":
        degraded = True

    lex_max = max((float(h.get("score", 0)) for h in lex_hits), default=0.0) or 1.0
    vec_max = max((float(h.get("score", 0)) for h in vec_hits), default=0.0) or 1.0

    merged: dict[str, dict[str, Any]] = {}

    for hit in lex_hits:
        path = hit["path"]
        if project and not path.startswith(project):
            continue
        norm = float(hit["score"]) / lex_max if lex_max else 0.0
        merged[path.lower()] = {
            "path": path,
            "lex": norm,
            "vec": 0.0,
            "snippet": hit.get("snippet", ""),
            "section_address": hit.get("section_address"),
            "source": "muse-lexical",
        }

    for hit in vec_hits:
        path = hit["path"]
        if project and not path.startswith(project):
            continue
        key = path.lower()
        norm = float(hit["score"]) / vec_max if vec_max else 0.0
        if key in merged:
            merged[key]["vec"] = norm
            merged[key]["source"] = "both"
            if hit.get("snippet") and not merged[key].get("snippet"):
                merged[key]["snippet"] = hit["snippet"]
        else:
            merged[key] = {
                "path": path,
                "lex": 0.0,
                "vec": norm,
                "snippet": hit.get("snippet", ""),
                "section_address": None,
                "source": "knowtation-vector",
            }

    # Case-insensitive basename dedupe — first walk-order path wins (lex order first).
    seen_basenames: dict[str, str] = {}
    dedupe_keys: list[str] = []
    for key in sorted(merged.keys(), key=lambda k: merged[k]["path"]):
        entry = merged[key]
        base = pathlib.PurePosixPath(entry["path"]).name.lower()
        if base in seen_basenames and seen_basenames[base] != entry["path"]:
            deduped += 1
            continue
        seen_basenames.setdefault(base, entry["path"])
        dedupe_keys.append(key)

    ranked: list[dict[str, Any]] = []
    for key in dedupe_keys:
        entry = merged[key]
        combined = W_LEX * entry["lex"] + W_VEC * entry["vec"]
        ranked.append({**entry, "combined": combined, "key": key})

    ranked.sort(
        key=lambda e: (
            -e["combined"],
            -_head_mtime(root, e["path"]),
            e["path"],
        )
    )

    results: list[dict[str, Any]] = []
    for entry in ranked[:top_k]:
        path = entry["path"]
        try:
            raw = (root / path).read_bytes()
            fm_dict, _ = _note_body_and_frontmatter(raw)
            frontmatter = {
                "project": fm_dict.get("project"),
                "tags": fm_dict.get("tags"),
                "source_type": fm_dict.get("source_type") or fm_dict.get("source"),
            }
        except OSError:
            frontmatter = {"project": None, "tags": None, "source_type": None}

        results.append(
            {
                "path": path,
                "score": round(entry["combined"], 6),
                "source": entry["source"],
                "snippet": entry.get("snippet", ""),
                "section_address": entry.get("section_address"),
                "frontmatter": frontmatter,
            }
        )

    if not results and lexical_status == "error" and vector_status != "ok":
        degraded = True

    return {
        "query": query,
        "mode": mode,
        "results": results,
        "index_freshness": index_freshness,
        "degraded": degraded,
        "deduped": deduped,
        "backends": {"lexical": lexical_status, "vector": vector_status},
    }


__all__ = [
    "W_LEX",
    "W_VEC",
    "VECTOR_TIMEOUT_SECONDS",
    "hybrid_search",
    "lexical_search",
    "vector_search",
]