"""Wiki keyword 召回 + graph BFS 扩展。"""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import Any

from niannian.brain.cerebellum.wiki_store import WikiStore


@dataclass
class WikiSearchHit:
    page_id: str
    title: str
    summary: str
    score: float
    source: str = "keyword"
    links: list[str] = field(default_factory=list)
    atom_ids: list[str] = field(default_factory=list)


def _wiki_recall_config() -> dict[str, Any]:
    from niannian.core.config import load_config

    wiki = (load_config().get("cerebellum") or {}).get("wiki") or {}
    recall = wiki.get("recall") or {}
    return recall if isinstance(recall, dict) else {}


class WikiSearch:
    def __init__(
        self,
        store,
        *,
        graph_bfs_enabled: bool | None = None,
        graph_bfs_depth: int | None = None,
    ) -> None:
        self._wiki = WikiStore(store)
        recall = _wiki_recall_config()
        if graph_bfs_enabled is None:
            graph_bfs_enabled = bool(recall.get("graph_bfs_enabled", True))
        if graph_bfs_depth is None:
            graph_bfs_depth = int(recall.get("graph_bfs_depth", 2))
        self._bfs_enabled = graph_bfs_enabled
        self._bfs_depth = max(0, graph_bfs_depth)

    def graph_bfs(
        self,
        from_ids: list[str],
        *,
        depth: int | None = None,
        namespace: str = "subject",
        object_id: str | None = None,
    ) -> list[str]:
        """从 seed 节点 BFS 扩 depth 跳；仅返回可见 namespace 内的 page_id。"""
        max_depth = self._bfs_depth if depth is None else max(0, depth)
        pages = self._wiki.pages_map()
        visible = self._wiki.filter_pages(pages, namespace=namespace, object_id=object_id)
        visible_ids = set(visible.keys())

        seeds = [pid for pid in from_ids if pid in visible_ids]
        visited: set[str] = set(seeds)
        frontier = list(seeds)
        expanded: list[str] = []

        for _ in range(max_depth):
            next_frontier: list[str] = []
            for pid in frontier:
                for nb in self._wiki.neighbors(pid):
                    if nb not in visible_ids or nb in visited:
                        continue
                    visited.add(nb)
                    next_frontier.append(nb)
                    expanded.append(nb)
            frontier = next_frontier
            if not frontier:
                break
        return expanded

    def _keyword_hits(self, query: str, pages: dict[str, Any]) -> list[WikiSearchHit]:
        q = query.lower()
        tokens = set(re.findall(r"[\w\u4e00-\u9fff]+", q))
        hits: list[WikiSearchHit] = []
        for page_id, page in pages.items():
            if not isinstance(page, dict):
                continue
            title = str(page.get("title") or page_id)
            links_raw = page.get("links") or []
            links = [str(x).removeprefix("ent:") for x in links_raw]
            blob = f"{title} {page.get('summary', '')} {' '.join(links)}".lower()
            score = float(sum(1 for t in tokens if t in blob))
            if title.lower() in q or page_id.lower() in q:
                score += 2.0
            if score <= 0:
                continue
            hits.append(
                WikiSearchHit(
                    page_id=page_id,
                    title=title,
                    summary=str(page.get("summary") or ""),
                    score=score,
                    source="keyword",
                    links=links,
                    atom_ids=list(page.get("atom_ids") or []),
                )
            )
        hits.sort(key=lambda h: h.score, reverse=True)
        return hits

    def search(
        self,
        query: str,
        *,
        namespace: str = "subject",
        object_id: str | None = None,
        limit: int = 20,
    ) -> list[WikiSearchHit]:
        pages = self._wiki.pages_map()
        filtered = self._wiki.filter_pages(pages, namespace=namespace, object_id=object_id)
        keyword_hits = self._keyword_hits(query, filtered)
        merged: dict[str, WikiSearchHit] = {h.page_id: h for h in keyword_hits}

        if self._bfs_enabled and keyword_hits:
            seed_ids = [h.page_id for h in keyword_hits]
            for pid in self.graph_bfs(
                seed_ids,
                namespace=namespace,
                object_id=object_id,
            ):
                if pid in merged:
                    continue
                page = filtered.get(pid)
                if not isinstance(page, dict):
                    continue
                links_raw = page.get("links") or []
                merged[pid] = WikiSearchHit(
                    page_id=pid,
                    title=str(page.get("title") or pid),
                    summary=str(page.get("summary") or ""),
                    score=0.5,
                    source="bfs",
                    links=[str(x).removeprefix("ent:") for x in links_raw],
                    atom_ids=list(page.get("atom_ids") or []),
                )

        hits = sorted(merged.values(), key=lambda h: h.score, reverse=True)
        return hits[:limit]
