"""大仓库：对已落盘的生成树做过滤、子树、统计，避免整份 JSON 塞进对话。"""

from __future__ import annotations

import json
import re
from collections import defaultdict
from pathlib import Path
from typing import Any


def load_tree_file(path: Path) -> dict[str, Any]:
    if not path.is_file():
        return {"ok": False, "error": "tree file not found"}
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as e:
        return {"ok": False, "error": str(e)}
    if not isinstance(data, dict):
        return {"ok": False, "error": "invalid tree json"}
    nodes = data.get("nodes")
    if not isinstance(nodes, list):
        return {"ok": False, "error": "missing nodes array"}
    return {"ok": True, "data": data, "nodes": nodes}


def _children_map(nodes: list[dict[str, Any]]) -> dict[str | None, list[dict[str, Any]]]:
    m: dict[str | None, list[dict[str, Any]]] = defaultdict(list)
    for n in nodes:
        pid = n.get("parent_id")
        if pid is not None and not isinstance(pid, str):
            pid = str(pid) if pid is not None else None
        m[pid].append(n)
    for k in m:
        m[k].sort(key=lambda x: (str(x.get("file", "")), int(x.get("line_start") or 0)))
    return m


def _collect_subtree_ids(
    cmap: dict[str | None, list[dict[str, Any]]],
    root_id: str,
    max_depth: int,
) -> set[str]:
    out: set[str] = {root_id}
    frontier = [(root_id, 0)]
    while frontier:
        nid, d = frontier.pop()
        if d >= max_depth:
            continue
        for ch in cmap.get(nid, []):
            cid = ch.get("id")
            if isinstance(cid, str) and cid not in out:
                out.add(cid)
                frontier.append((cid, d + 1))
    return out


def query_tree(
    path: Path,
    *,
    action: str,
    parent_id: str | None = None,
    subtree_depth: int = 8,
    qname_contains: str | None = None,
    file_contains: str | None = None,
    kinds: list[str] | None = None,
    source_langs: list[str] | None = None,
    limit: int = 200,
    offset: int = 0,
    include_docstrings: bool = False,
) -> dict[str, Any]:
    loaded = load_tree_file(path)
    if not loaded.get("ok"):
        return loaded
    data = loaded["data"]
    nodes: list[dict[str, Any]] = loaded["nodes"]
    meta = {
        "index_revision": data.get("index_revision"),
        "topic_slug": data.get("topic_slug"),
        "schema": data.get("schema"),
        "scan_root": data.get("root"),
        "files_scanned": data.get("files_scanned"),
        "total_nodes": len(nodes),
        "tree_path": str(path.resolve()),
    }

    if action == "stats":
        by_kind: dict[str, int] = defaultdict(int)
        by_lang: dict[str, int] = defaultdict(int)
        by_file: dict[str, int] = defaultdict(int)
        for n in nodes:
            by_kind[str(n.get("kind") or "?")] += 1
            by_lang[str(n.get("source_lang") or "python")] += 1
            fp = str(n.get("file") or "")
            by_file[fp] += 1
        top_files = sorted(by_file.items(), key=lambda x: -x[1])[:40]
        return {
            "ok": True,
            "action": "stats",
            "meta": meta,
            "by_kind": dict(by_kind),
            "by_source_lang": dict(by_lang),
            "distinct_files": len(by_file),
            "top_files_by_node_count": [{"file": a, "nodes": b} for a, b in top_files],
        }

    cmap = _children_map(nodes)
    kind_set = {k.lower() for k in kinds} if kinds else None
    lang_set = {k.lower() for k in source_langs} if source_langs else None
    qn_sub = (qname_contains or "").lower()
    fc_sub = (file_contains or "").lower()

    def _filter_base(seq: list[dict[str, Any]]) -> list[dict[str, Any]]:
        out: list[dict[str, Any]] = []
        for n in seq:
            if kind_set and str(n.get("kind") or "").lower() not in kind_set:
                continue
            if lang_set:
                sl = str(n.get("source_lang") or "python").lower()
                if sl not in lang_set:
                    continue
            if qn_sub and qn_sub not in str(n.get("qname") or "").lower():
                continue
            if fc_sub and fc_sub not in str(n.get("file") or "").lower():
                continue
            out.append(n)
        return out

    def _shrink(n: dict[str, Any]) -> dict[str, Any]:
        keys = [
            "id",
            "parent_id",
            "kind",
            "name",
            "qname",
            "file",
            "line_start",
            "line_end",
            "source_lang",
            "parser",
        ]
        r = {k: n.get(k) for k in keys}
        if include_docstrings:
            r["docstring_preview"] = n.get("docstring_preview")
            r["ai_description_slot"] = n.get("ai_description_slot")
        return r

    if action == "children":
        if not parent_id:
            return {"ok": False, "error": "children action requires parent_id"}
        raw = cmap.get(parent_id, [])
        filt = _filter_base(raw)
        slice_f = filt[offset : offset + max(1, min(500, limit))]
        return {
            "ok": True,
            "action": "children",
            "meta": meta,
            "parent_id": parent_id,
            "returned": len(slice_f),
            "total_matched": len(filt),
            "nodes": [_shrink(x) for x in slice_f],
        }

    if action == "subtree":
        if not parent_id:
            return {"ok": False, "error": "subtree action requires parent_id"}
        depth = max(1, min(64, subtree_depth))
        ids = _collect_subtree_ids(cmap, parent_id, depth)
        sub = [n for n in nodes if n.get("id") in ids]
        filt = _filter_base(sub)
        slice_f = filt[offset : offset + max(1, min(500, limit))]
        return {
            "ok": True,
            "action": "subtree",
            "meta": meta,
            "root_id": parent_id,
            "subtree_depth_limit": depth,
            "returned": len(slice_f),
            "total_matched": len(filt),
            "nodes": [_shrink(x) for x in slice_f],
        }

    if action == "filter":
        filt = _filter_base(nodes)
        slice_f = filt[offset : offset + max(1, min(500, limit))]
        return {
            "ok": True,
            "action": "filter",
            "meta": meta,
            "returned": len(slice_f),
            "total_matched": len(filt),
            "offset": offset,
            "limit": limit,
            "nodes": [_shrink(x) for x in slice_f],
        }

    return {"ok": False, "error": f"unknown action: {action}"}


def query_tree_regex(
    path: Path,
    *,
    qname_regex: str,
    limit: int = 100,
) -> dict[str, Any]:
    try:
        rx = re.compile(qname_regex, re.I)
    except re.error as e:
        return {"ok": False, "error": f"invalid regex: {e}"}
    loaded = load_tree_file(path)
    if not loaded.get("ok"):
        return loaded
    data = loaded["data"]
    nodes = loaded["nodes"]
    meta = {
        "index_revision": data.get("index_revision"),
        "topic_slug": data.get("topic_slug"),
        "total_nodes": len(nodes),
        "tree_path": str(path.resolve()),
    }
    out = [n for n in nodes if rx.search(str(n.get("qname") or ""))]
    out = out[: max(1, min(300, limit))]
    return {
        "ok": True,
        "action": "regex",
        "meta": meta,
        "returned": len(out),
        "nodes": [
            {
                "id": n.get("id"),
                "parent_id": n.get("parent_id"),
                "kind": n.get("kind"),
                "qname": n.get("qname"),
                "file": n.get("file"),
                "line_start": n.get("line_start"),
                "line_end": n.get("line_end"),
            }
            for n in out
        ],
    }
