"""Hierarchy-centric mini GraphRAG index builder and lexical search."""

from __future__ import annotations

import json
import re
from collections import defaultdict
from pathlib import Path
from typing import Any, Dict, List, Set

from ..service import HierarchyService

STOPWORDS: Set[str] = {
    "the", "a", "an", "to", "for", "and", "or", "of", "in", "on", "with", "by", "is", "are",
}


def _tokenize(text: str) -> List[str]:
    parts = re.findall(r"[a-zA-Z0-9_]+", (text or "").lower())
    return [t for t in parts if len(t) > 1 and t not in STOPWORDS]


def build_hierarchy_rag(
    project_id: str,
    data_dir: str = "data",
    index_path: str = "data/hierarchy_rag/index.json",
) -> Dict[str, Any]:
    service = HierarchyService(data_dir=data_dir)
    project = service.get_project(project_id)
    if not project:
        return {"error": f"Project '{project_id}' not found"}
    hierarchies = service.list_hierarchies(project_id)

    nodes: Dict[str, Dict[str, Any]] = {}
    edges: Dict[str, List[Dict[str, str]]] = defaultdict(list)
    inv: Dict[str, Set[str]] = defaultdict(set)

    for h in hierarchies:
        hid = h.get("hierarchy_id", "")
        nid = f"hierarchy_node:{project_id}:{hid}"
        levels = h.get("hierarchy_level", {}) or {}
        level_depth = sum(1 for i in range(1, 16) if levels.get(f"level_{i}"))
        node = {
            "id": nid,
            "type": "hierarchy_node",
            "project_id": project_id,
            "name": h.get("hierarchy_name", ""),
            "hierarchy_id": hid,
            "parent_id": h.get("parent_id", ""),
            "level_depth": level_depth,
            "mapping_count": len(h.get("mapping", [])),
            "has_formula": bool((h.get("formula_config", {}) or {}).get("formula_group")),
        }
        nodes[nid] = node

        parent = h.get("parent_id")
        if parent:
            edges[nid].append({"to": f"hierarchy_node:{project_id}:{parent}", "type": "rolls_up_to"})

        if h.get("dimension_props"):
            edges[nid].append({"to": f"kpi:{project_id}:dimension_quality", "type": "classified_as"})

        for m in h.get("mapping", []) or []:
            table = m.get("source_table", "")
            col = m.get("source_column", "")
            if not table:
                continue
            fact_id = f"fact:{project_id}:{table}"
            if fact_id not in nodes:
                nodes[fact_id] = {
                    "id": fact_id,
                    "type": "fact",
                    "project_id": project_id,
                    "name": table,
                }
            edges[nid].append({"to": fact_id, "type": "depends_on"})
            if col:
                meas_id = f"measure:{project_id}:{table}.{col}"
                if meas_id not in nodes:
                    nodes[meas_id] = {
                        "id": meas_id,
                        "type": "measure",
                        "project_id": project_id,
                        "name": f"{table}.{col}",
                    }
                edges[fact_id].append({"to": meas_id, "type": "drives"})

    # Build lexical inverted index.
    for nid, n in nodes.items():
        blob = " ".join(str(n.get(k, "")) for k in ["id", "type", "name", "project_id", "hierarchy_id"])
        for t in _tokenize(blob):
            inv[t].add(nid)

    payload = {
        "schema_version": "1.0",
        "project_id": project_id,
        "nodes": nodes,
        "edges": {k: v for k, v in edges.items()},
        "inverted_index": {k: sorted(list(v)) for k, v in inv.items()},
        "stats": {
            "node_count": len(nodes),
            "edge_count": sum(len(v) for v in edges.values()),
            "token_count": len(inv),
        },
    }
    out = Path(index_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    return {"status": "ok", "index_path": index_path, "stats": payload["stats"]}


def search_hierarchy_rag(index_path: str, query: str, k: int = 10) -> Dict[str, Any]:
    p = Path(index_path)
    if not p.exists():
        return {"error": f"Hierarchy RAG index not found at '{index_path}'"}
    index = json.loads(p.read_text(encoding="utf-8"))
    inv = index.get("inverted_index", {}) or {}
    nodes = index.get("nodes", {}) or {}

    scores: Dict[str, float] = defaultdict(float)
    q_tokens = _tokenize(query)
    for t in q_tokens:
        ids = inv.get(t, [])
        if not ids:
            continue
        idf = 1.0 / max(1, len(ids))
        for nid in ids:
            scores[nid] += 1.0 + idf

    ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)[: max(1, int(k))]
    results = []
    for nid, s in ranked:
        n = nodes.get(nid, {})
        results.append(
            {
                "id": nid,
                "score": round(s, 4),
                "type": n.get("type"),
                "name": n.get("name"),
                "project_id": n.get("project_id"),
            }
        )
    return {"status": "ok", "query": query, "count": len(results), "results": results}

