"""
llm_clusterer.py - Stage 4 (LLM variant): semantic co-location clustering.

An alternative to the deterministic clusterer.py. Sends the flattened
assertions (each with a stable id) to the model, which groups ones that refer
to the same real-world room/door/grid/tag/equipment -- catching fuzzy matches
the exact-key clusterer misses ("Room 124" vs "Mgr Office 124"). The returned
assertion_ids are resolved back to full assertions and emitted in the SAME
internal cluster shape conflict_checker/report consume, so it is a drop-in swap.

Selected via config.CLUSTERER == "llm". On failure returns []. Keeps clusters
with >= 2 member assertions (a single-assertion cluster gives the reasoner
nothing to compare).
"""

import json
from typing import Dict, List

from backend import config
from backend.pipeline._stage import call_stage, collect_list
from backend.prompts import CLUSTER_SYSTEM_PROMPT, CLUSTER_USER_INSTRUCTION

MAX_ASSERTIONS_PER_CLUSTER = 24


def _flatten(sheets: List[Dict]) -> List[Dict]:
    flat: List[Dict] = []
    for s in sheets:
        for a in s.get("assertions", []):
            flat.append({**a,
                         "discipline": s.get("discipline", "Unknown"),
                         "sheet_number": s.get("sheet_number"),
                         "page_number": s.get("page_number")})
    return flat


def _payload(flat: List[Dict]) -> List[Dict]:
    return [{
        "assertion_id": a["id"],
        "discipline": a["discipline"],
        "sheet_number": a.get("sheet_number"),
        "attribute": a.get("attribute"),
        "value": a.get("value"),
        "location_key": a.get("location_key"),
        "source_text": a.get("source_text"),
    } for a in flat]


def _location(pk) -> str:
    if not isinstance(pk, dict):
        return "Unspecified"
    lvl = f" / {pk['level']}" if pk.get("level") else ""
    for key, label in (("room", "Room"), ("grid", "Grid"), ("tag", "Tag")):
        if pk.get(key):
            return f"{label} {pk[key]}{lvl}"
    return pk.get("plan_zone") or pk.get("room_name") or "Unspecified"


def cluster_by_location_llm(sheets: List[Dict]) -> List[Dict]:
    flat = _flatten(sheets)
    if not flat:
        return []
    by_id = {a["id"]: a for a in flat}

    parsed = call_stage(
        CLUSTER_SYSTEM_PROMPT,
        CLUSTER_USER_INSTRUCTION,
        subs={"normalized_assertions": json.dumps(_payload(flat), ensure_ascii=True)},
        max_tokens=config.CLUSTER_MAX_TOKENS,
    )
    raw = collect_list(parsed, "clusters")

    clusters: List[Dict] = []
    for c in raw:
        member_ids = [i for i in (c.get("assertion_ids") or []) if i in by_id]
        members = [by_id[i] for i in member_ids]
        disciplines = sorted({m["discipline"] for m in members})
        # Need at least two facts to compare (cross-discipline, or schedule vs plan).
        if len(members) < 2:
            continue
        members = members[:MAX_ASSERTIONS_PER_CLUSTER]
        clusters.append({
            "key": str(c.get("cluster_id") or _location(c.get("primary_location_key"))),
            "location": _location(c.get("primary_location_key")),
            "disciplines": disciplines,
            "page_numbers": sorted({m["page_number"] for m in members if m.get("page_number")}),
            "sheets": sorted({m["sheet_number"] for m in members if m.get("sheet_number")}),
            "assertions": members,
            "kind": "llm",
        })

    clusters.sort(key=lambda c: (-len(c["disciplines"]), -len(c["assertions"]), c["key"]))
    if len(clusters) > config.CLUSTER_MAX:
        print(f"[Cluster/LLM] capping {len(clusters)} -> {config.CLUSTER_MAX}")
        clusters = clusters[:config.CLUSTER_MAX]
    print(f"[Cluster/LLM] {len(clusters)} clusters (from {len(raw)} returned)")
    return clusters