"""
ZIA Kuzu Graph Adapter — tenant-namespaced graph operations over Kuzu.

Key design decisions:
- Tenant isolation via table name prefixing: tenant__{id}__asset, etc.
- No raw Cypher/graph queries escape the adapter — all operations go through typed methods.
- Phase 1: embedded Kuzu (file-based). Phase 2: swap to Memgraph (server-mode) behind same interface.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Optional

try:
    import kuzu  # type: ignore[import-untyped]
except ImportError:
    kuzu = None  # type: ignore[assignment]

from zia.sif.schema.nodes import (
    AIModelNode,
    AssetNode,
    ControlNode,
    IdentityNode,
    RiskNode,
    SIFNode,
    VendorNode,
    VulnerabilityNode,
)


# Map node type names to their Pydantic classes
NODE_TYPE_MAP: dict[str, type[SIFNode]] = {
    "asset": AssetNode,
    "identity": IdentityNode,
    "vulnerability": VulnerabilityNode,
    "control": ControlNode,
    "risk": RiskNode,
    "vendor": VendorNode,
    "ai_model": AIModelNode,
}


class KuzuAdapter:
    """
    Tenant-namespaced graph adapter over Kuzu (embedded graph DB).

    Every operation requires a tenant_id. Table names are prefixed with the
    tenant namespace so data from different tenants never mixes at the storage level.

    Usage:
        adapter = KuzuAdapter(db_path="/tmp/sif.db")
        adapter.initialize_schema(tenant_id="acme")
        adapter.upsert_node(tenant_id="acme", node=AssetNode(...))
        nodes = adapter.get_nodes(tenant_id="acme", node_type="asset")
    """

    def __init__(self, db_path: str | Path = "./sif_graph") -> None:
        self.db_path = Path(db_path)
        self.db_path.mkdir(parents=True, exist_ok=True)
        self._db = kuzu.Database(str(self.db_path))
        self._conn = kuzu.Connection(self._db)

    def _ns(self, tenant_id: str, node_type: str) -> str:
        """Return the namespaced table name for a tenant + node type."""
        safe = tenant_id.replace("-", "_").lower()
        return f"t_{safe}_{node_type}"

    def initialize_schema(self, tenant_id: str) -> None:
        """
        Create Kuzu node tables for a tenant if they don't already exist.

        Called once on tenant registration. Safe to call multiple times (idempotent).
        """
        for node_type in NODE_TYPE_MAP:
            table_name = self._ns(tenant_id, node_type)
            # All nodes share a common set of base columns + a JSON blob for extra fields
            self._conn.execute(f"""
                CREATE NODE TABLE IF NOT EXISTS {table_name} (
                    node_id STRING,
                    valid_from STRING,
                    valid_to STRING,
                    confidence DOUBLE,
                    source STRING,
                    data STRING,
                    PRIMARY KEY (node_id)
                )
            """)

        # Edge tables
        for edge_def in [
            ("identity_accesses_asset", "identity", "asset"),
            ("asset_has_vulnerability", "asset", "vulnerability"),
            ("control_mitigates_vuln", "control", "vulnerability"),
            ("vendor_supplies_asset", "vendor", "asset"),
            ("risk_impacts_asset", "risk", "asset"),
            ("asset_communicates_with", "asset", "asset"),
        ]:
            edge_table, from_type, to_type = edge_def
            edge_name = self._ns(tenant_id, edge_table)
            from_table = self._ns(tenant_id, from_type)
            to_table = self._ns(tenant_id, to_type)
            self._conn.execute(f"""
                CREATE REL TABLE IF NOT EXISTS {edge_name} (
                    FROM {from_table} TO {to_table},
                    valid_from STRING,
                    confidence DOUBLE,
                    source STRING,
                    data STRING
                )
            """)

    def upsert_node(self, tenant_id: str, node: SIFNode) -> None:
        """
        Insert or update a node in the graph.
        Determines node type from the Pydantic class name.
        """
        node_type = _get_node_type_key(node)
        table_name = self._ns(tenant_id, node_type)

        data = node.model_dump(mode="json")
        node_id = data.pop("node_id")
        valid_from = str(data.pop("valid_from", ""))
        valid_to = str(data.pop("valid_to", "") or "")
        confidence = float(data.pop("confidence", 1.0))
        source = str(data.pop("source", ""))
        extra_json = json.dumps(data)

        # Kuzu MERGE equivalent: delete + insert (Kuzu doesn't support UPSERT yet)
        self._conn.execute(
            f"MATCH (n:{table_name} {{node_id: $node_id}}) DELETE n",
            parameters={"node_id": node_id},
        )
        self._conn.execute(
            f"""
            CREATE (:{table_name} {{
                node_id: $node_id,
                valid_from: $valid_from,
                valid_to: $valid_to,
                confidence: $confidence,
                source: $source,
                data: $data
            }})
            """,
            parameters={
                "node_id": node_id,
                "valid_from": valid_from,
                "valid_to": valid_to,
                "confidence": confidence,
                "source": source,
                "data": extra_json,
            },
        )

    def get_node(
        self, tenant_id: str, node_type: str, node_id: str
    ) -> Optional[dict[str, Any]]:
        """Retrieve a single node by ID. Returns None if not found."""
        table_name = self._ns(tenant_id, node_type)
        result = self._conn.execute(
            f"MATCH (n:{table_name} {{node_id: $node_id}}) RETURN n",
            parameters={"node_id": node_id},
        )
        rows = result.get_as_df() if result else None
        if rows is None or rows.empty:
            return None
        row = rows.iloc[0]["n"]
        return _deserialise_row(row)

    def get_nodes(
        self, tenant_id: str, node_type: str, filters: Optional[dict[str, Any]] = None
    ) -> list[dict[str, Any]]:
        """Retrieve all nodes of a given type for a tenant, with optional dict filters."""
        table_name = self._ns(tenant_id, node_type)
        result = self._conn.execute(f"MATCH (n:{table_name}) RETURN n")
        if result is None:
            return []
        df = result.get_as_df()
        if df.empty:
            return []
        nodes = [_deserialise_row(row["n"]) for _, row in df.iterrows()]
        # Apply in-memory filters (Phase 1 simplicity)
        if filters:
            nodes = [
                n for n in nodes
                if all(n.get(k) == v for k, v in filters.items())
            ]
        return nodes

    def close(self) -> None:
        del self._conn
        del self._db


def _get_node_type_key(node: SIFNode) -> str:
    """Map a Pydantic node instance to its string type key."""
    reverse = {v: k for k, v in NODE_TYPE_MAP.items()}
    key = reverse.get(type(node))
    if key is None:
        raise ValueError(f"Unknown node type: {type(node).__name__}. Register it in NODE_TYPE_MAP.")
    return key


def _deserialise_row(row: Any) -> dict[str, Any]:
    """Flatten a Kuzu row dict + its JSON 'data' blob into one dict."""
    result = dict(row)
    extra = result.pop("data", "{}")
    result.update(json.loads(extra))
    return result
