"""Snowflake graph adapter (Phase 6)."""
from __future__ import annotations

import json
import os
from typing import Iterable, List

from ..types import GraphNode, GraphEdge


class SnowflakeGraphAdapter:
    """Minimal Snowflake graph adapter using table storage."""

    def __init__(self, database: str, schema: str):
        self.database = database
        self.schema = schema
        try:
            import snowflake.connector
        except Exception as exc:
            raise RuntimeError("snowflake-connector-python not installed") from exc
        self._connector = snowflake.connector
        self._ensure_tables()

    def _connect(self):
        return self._connector.connect(
            account=os.getenv("SNOWFLAKE_ACCOUNT", ""),
            user=os.getenv("SNOWFLAKE_USER", ""),
            password=os.getenv("SNOWFLAKE_PASSWORD", ""),
            role=os.getenv("SNOWFLAKE_ROLE", ""),
            warehouse=os.getenv("SNOWFLAKE_WAREHOUSE", ""),
            database=self.database,
            schema=self.schema,
        )

    def _ensure_tables(self) -> None:
        with self._connect() as conn:
            cur = conn.cursor()
            cur.execute(
                "CREATE TABLE IF NOT EXISTS KB_GRAPH_NODES ("
                "ID STRING PRIMARY KEY, "
                "TYPE STRING, "
                "PROPERTIES VARIANT)"
            )
            cur.execute(
                "CREATE TABLE IF NOT EXISTS KB_GRAPH_EDGES ("
                "SOURCE STRING, "
                "TARGET STRING, "
                "EDGE_TYPE STRING, "
                "PROPERTIES VARIANT, "
                "PRIMARY KEY (SOURCE, TARGET, EDGE_TYPE))"
            )
            cur.close()

    def upsert_nodes(self, nodes: Iterable[GraphNode]) -> int:
        with self._connect() as conn:
            cur = conn.cursor()
            count = 0
            for node in nodes:
                cur.execute(
                    "MERGE INTO KB_GRAPH_NODES t "
                    "USING (SELECT %s AS ID, %s AS TYPE, PARSE_JSON(%s) AS PROPERTIES) s "
                    "ON t.ID = s.ID "
                    "WHEN MATCHED THEN UPDATE SET TYPE = s.TYPE, PROPERTIES = s.PROPERTIES "
                    "WHEN NOT MATCHED THEN INSERT (ID, TYPE, PROPERTIES) VALUES (s.ID, s.TYPE, s.PROPERTIES)",
                    (node.id, node.type, json.dumps(node.properties)),
                )
                count += 1
            cur.close()
        return count

    def upsert_edges(self, edges: Iterable[GraphEdge]) -> int:
        with self._connect() as conn:
            cur = conn.cursor()
            count = 0
            for edge in edges:
                cur.execute(
                    "MERGE INTO KB_GRAPH_EDGES t "
                    "USING (SELECT %s AS SOURCE, %s AS TARGET, %s AS EDGE_TYPE, PARSE_JSON(%s) AS PROPERTIES) s "
                    "ON t.SOURCE = s.SOURCE AND t.TARGET = s.TARGET AND t.EDGE_TYPE = s.EDGE_TYPE "
                    "WHEN MATCHED THEN UPDATE SET PROPERTIES = s.PROPERTIES "
                    "WHEN NOT MATCHED THEN INSERT (SOURCE, TARGET, EDGE_TYPE, PROPERTIES) "
                    "VALUES (s.SOURCE, s.TARGET, s.EDGE_TYPE, s.PROPERTIES)",
                    (edge.source, edge.target, edge.type, json.dumps(edge.properties)),
                )
                count += 1
            cur.close()
        return count

    def persist_run_graph(self, run_id: str, nodes: List[GraphNode], edges: List[GraphEdge]) -> int:
        """Persist a complete run graph as VARIANT rows in a per-run table.

        Each run's graph is stored as a JSON VARIANT row keyed by RUN_ID
        in KB_RUN_GRAPHS for historical analysis and auditability.

        Args:
            run_id: Unique workflow run identifier.
            nodes: List of GraphNode objects from this run.
            edges: List of GraphEdge objects from this run.

        Returns:
            Total number of items persisted (nodes + edges).
        """
        with self._connect() as conn:
            cur = conn.cursor()
            cur.execute(
                "CREATE TABLE IF NOT EXISTS KB_RUN_GRAPHS ("
                "RUN_ID STRING, "
                "ITEM_TYPE STRING, "
                "ITEM_ID STRING, "
                "DATA VARIANT, "
                "CREATED_AT TIMESTAMP_NTZ DEFAULT CURRENT_TIMESTAMP())"
            )
            count = 0
            for node in nodes:
                cur.execute(
                    "INSERT INTO KB_RUN_GRAPHS (RUN_ID, ITEM_TYPE, ITEM_ID, DATA) "
                    "SELECT %s, 'node', %s, PARSE_JSON(%s)",
                    (run_id, node.id, json.dumps({"type": node.type, "properties": node.properties})),
                )
                count += 1
            for edge in edges:
                edge_id = f"{edge.source}|{edge.target}|{edge.type}"
                cur.execute(
                    "INSERT INTO KB_RUN_GRAPHS (RUN_ID, ITEM_TYPE, ITEM_ID, DATA) "
                    "SELECT %s, 'edge', %s, PARSE_JSON(%s)",
                    (run_id, edge_id, json.dumps({"source": edge.source, "target": edge.target, "type": edge.type, "properties": edge.properties})),
                )
                count += 1
            cur.close()
        # Also upsert into main tables
        self.upsert_nodes(nodes)
        self.upsert_edges(edges)
        return count

    def get_edges(self) -> List[GraphEdge]:
        edges: List[GraphEdge] = []
        with self._connect() as conn:
            cur = conn.cursor()
            cur.execute("SELECT SOURCE, TARGET, EDGE_TYPE, PROPERTIES FROM KB_GRAPH_EDGES")
            for row in cur.fetchall():
                props = row[3]
                if isinstance(props, str):
                    try:
                        props = json.loads(props)
                    except json.JSONDecodeError:
                        props = {}
                edges.append(GraphEdge(source=row[0], target=row[1], type=row[2], properties=props or {}))
            cur.close()
        return edges
