"""/graph/explore — neighbourhood exploration filtered by predicate sets per mode (API 1.1). `/entities/{slug}/graph` (v1) is unchanged."""
from __future__ import annotations

from typing import Any

from fastapi import APIRouter, Query, Request

from aiatlas.api.common import ApiError, cached, resolve_entity
from aiatlas.db import connection, fetch_all

router = APIRouter(prefix="/api/v1/graph", tags=["graph"])

MODES: dict[str, tuple[str, ...]] = {
    "lineage": ("artifact_of", "quantized_from", "fine_tuned_from", "distilled_from", "merged_from", "derived_from", "superseded_by", "member_of_family", "develops"),
    "research": ("authored", "described_by", "published_by", "uses_dataset", "evaluates_on", "evaluated_on"),
    "company": ("develops", "owns", "operates", "manufactures", "member_of_family"),
    "benchmark": ("evaluated_on", "evaluates_on", "variant_of"),
    "dataset": ("uses_dataset", "derived_from", "subset_of"),
    "provider": ("available_through",),
    "hardware": ("runs_on", "uses", "manufactures"),
}
KEY_ATTRS = ("parameter_count", "context_length", "release_date", "openness", "published_at", "category", "memory_gb", "kind")


@router.get("/explore")
@cached(300)
async def explore(request: Request, node: str = Query(...), mode: str = Query("lineage"), depth: int = Query(1, ge=1, le=2), limit: int = Query(150, ge=2, le=150)) -> dict[str, Any]:
    if mode not in MODES:
        raise ApiError(400, f"mode must be one of {', '.join(MODES)}")
    preds = list(MODES[mode])
    async with connection() as conn:
        root_row = await resolve_entity(conn, node)
        root = root_row["id"]
        seen: dict[str, int] = {root: 0}
        frontier = [root]
        edges: dict[tuple[str, str, str], dict[str, Any]] = {}
        truncated = False
        for level in range(1, depth + 1):
            if not frontier:
                break
            rows = await fetch_all(conn, """select r.subject_id, r.predicate, r.object_id, r.attributes, r.tier from relations r
                                            where r.valid_to is null and r.predicate = any(cast(:preds as text[]))
                                              and (r.subject_id = any(cast(:ids as text[])) or r.object_id = any(cast(:ids as text[])))
                                            order by r.tier, r.observed_at desc limit :lim""", preds=preds, ids=frontier, lim=limit * 4)
            nxt: list[str] = []
            for r in rows:
                other = r["object_id"] if r["subject_id"] in seen else r["subject_id"]
                if other not in seen:
                    if len(seen) >= limit:
                        truncated = True
                        continue
                    seen[other] = level
                    nxt.append(other)
                edges[(r["subject_id"], r["object_id"], r["predicate"])] = {"source": r["subject_id"], "target": r["object_id"], "predicate": r["predicate"], "attributes": r["attributes"] or {}, "tier": r["tier"]}
            frontier = nxt
        # organisation edges for the company mode come from entities.organization_id as well
        if mode == "company":
            org_rows = await fetch_all(conn, "select id, organization_id from entities where organization_id is not null and merged_into is null and (id = any(cast(:ids as text[])) or organization_id = any(cast(:ids as text[]))) limit :lim",
                                       ids=list(seen), lim=limit * 4)
            for r in org_rows:
                for x in (r["id"], r["organization_id"]):
                    if x not in seen:
                        if len(seen) >= limit:
                            truncated = True
                            break
                        seen[x] = depth
                edges.setdefault((r["organization_id"], r["id"], "develops"), {"source": r["organization_id"], "target": r["id"], "predicate": "develops", "attributes": {"from": "organization_id"}, "tier": None})
        nodes = await fetch_all(conn, """select e.id, e.slug, e.canonical_name, e.entity_type, e.attributes, e.merged_into, e.artifact_kind, o.canonical_name as org_name, o.slug as org_slug
                                         from entities e left join entities o on o.id = e.organization_id where e.id = any(cast(:ids as text[]))""", ids=list(seen))
    node_ids = {n["id"] for n in nodes}
    out_edges = [e for e in edges.values() if e["source"] in node_ids and e["target"] in node_ids]
    counts: dict[str, int] = {}
    for n in nodes:
        counts[n["entity_type"]] = counts.get(n["entity_type"], 0) + 1
    return {"root": root, "mode": mode, "depth": depth, "predicates": preds,
            "nodes": [{"id": n["id"], "slug": n["slug"], "name": n["canonical_name"], "entity_type": n["entity_type"], "org": n["org_name"], "org_slug": n["org_slug"], "level": seen.get(n["id"], 0),
                       "artifact_kind": n["artifact_kind"], "attributes": {k: (n["attributes"] or {}).get(k) for k in KEY_ATTRS if (n["attributes"] or {}).get(k) not in (None, "", [])}} for n in nodes],
            "edges": out_edges, "truncated": truncated, "counts": {"nodes": len(nodes), "edges": len(out_edges), "by_type": counts}}