"""Snowflake-backed persistence adapter (VARIANT table)."""

from __future__ import annotations

import json
import os
from typing import Any, Dict, List, Optional

from .adapter import PersistenceAdapter
from .snowflake_bulk import SnowflakeBulkLoader


class SnowflakePersistenceAdapter(PersistenceAdapter):
    """Persist workflow runs to a Snowflake VARIANT table via MERGE upserts."""

    TABLE_NAME = "PLATFORM_RUN_SUMMARY"

    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_table()

    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_table(self) -> None:
        with self._connect() as conn:
            cur = conn.cursor()
            cur.execute(
                f"CREATE TABLE IF NOT EXISTS {self.TABLE_NAME} ("
                "RUN_ID STRING PRIMARY KEY, "
                "SUMMARY VARIANT, "
                "CREATED_AT TIMESTAMP_NTZ DEFAULT CURRENT_TIMESTAMP())"
            )
            cur.close()

    def save_run(self, run_id: str, summary: Dict[str, Any]) -> None:
        payload = json.dumps(summary, default=str)
        with self._connect() as conn:
            cur = conn.cursor()
            # Use bulk staged path for large payloads (>500KB)
            if len(payload) > 512_000:
                loader = SnowflakeBulkLoader()
                loader.bulk_merge_variant(
                    cursor=cur,
                    table=self.TABLE_NAME,
                    id_col="RUN_ID",
                    variant_col="SUMMARY",
                    records=[{"RUN_ID": run_id, **summary}],
                )
            else:
                cur.execute(
                    f"MERGE INTO {self.TABLE_NAME} t "
                    "USING (SELECT %s AS RUN_ID, PARSE_JSON(%s) AS SUMMARY) s "
                    "ON t.RUN_ID = s.RUN_ID "
                    "WHEN MATCHED THEN UPDATE SET SUMMARY = s.SUMMARY "
                    "WHEN NOT MATCHED THEN INSERT (RUN_ID, SUMMARY) VALUES (s.RUN_ID, s.SUMMARY)",
                    (run_id, payload),
                )
            cur.close()

    def load_run(self, run_id: str) -> Optional[Dict[str, Any]]:
        with self._connect() as conn:
            cur = conn.cursor()
            cur.execute(
                f"SELECT SUMMARY FROM {self.TABLE_NAME} WHERE RUN_ID = %s", (run_id,)
            )
            row = cur.fetchone()
            cur.close()
            if not row:
                return None
            data = row[0]
            if isinstance(data, str):
                return json.loads(data)
            return data

    def save_artifact_meta(self, run_id: str, artifact: Dict[str, Any]) -> None:
        run = self.load_run(run_id)
        if run is None:
            run = {"run_id": run_id, "artifacts": []}
        arts = run.get("artifacts", [])
        arts.append(artifact)
        run["artifacts"] = arts
        self.save_run(run_id, run)

    def list_runs(self, limit: int = 20) -> List[Dict[str, Any]]:
        with self._connect() as conn:
            cur = conn.cursor()
            cur.execute(
                f"SELECT SUMMARY FROM {self.TABLE_NAME} "
                "ORDER BY CREATED_AT DESC LIMIT %s",
                (limit,),
            )
            rows = cur.fetchall()
            cur.close()
        results: List[Dict[str, Any]] = []
        for row in rows:
            data = row[0]
            if isinstance(data, str):
                data = json.loads(data)
            results.append(data)
        return results
