"""
Baseline loading strategies (OCP: open for extension via new loaders).

Each loader implements :class:`BaselineLoader` and knows how to load
baselines from a single source.  The :func:`create_baseline_loader`
factory selects the appropriate loader based on the configuration.

The actual baseline *data* for Snowflake comparison is read directly
from the stage by ``materialize_baseline_transient_table`` at comparison
time.  The loader only needs to discover which baselines exist and
their success/error status.
"""

from __future__ import annotations

import json
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, Optional

from snowflake.snowflake_data_validation.connector.connector_base import ConnectorBase

from test_runner.common.baseline import discover_baselines, load_baseline
from test_runner.common.models import BaselineData


# ---------------------------------------------------------------------------
# Abstract loader
# ---------------------------------------------------------------------------

class BaselineLoader(ABC):
    """Loads baselines from a single source."""

    @abstractmethod
    def load(self) -> list[BaselineData]:
        """Load all available baselines from this source."""

    @property
    @abstractmethod
    def source_label(self) -> str:
        """Human-readable label describing this baseline source."""


# ---------------------------------------------------------------------------
# Disk loader (offline / testing)
# ---------------------------------------------------------------------------

class DiskBaselineLoader(BaselineLoader):
    """Load baselines from local JSON files on disk."""

    def __init__(self, baseline_dir: Path) -> None:
        self._baseline_dir = baseline_dir

    def load(self) -> list[BaselineData]:
        baselines: list[BaselineData] = []
        for path in discover_baselines(self._baseline_dir):
            bl = load_baseline(path)
            if bl is not None:
                baselines.append(bl)
        return baselines

    @property
    def source_label(self) -> str:
        return f"{self._baseline_dir} (local)"


# ---------------------------------------------------------------------------
# Stage loader
# ---------------------------------------------------------------------------

def _parse_variant(raw: Any) -> Optional[dict[str, Any]]:
    """Parse a Snowflake VARIANT value into a Python dict.

    The connector may return VARIANT as a JSON string or as a dict
    depending on the driver version and result format.
    """
    if raw is None:
        return None
    if isinstance(raw, dict):
        return raw
    if isinstance(raw, str):
        try:
            parsed = json.loads(raw)
            return parsed if isinstance(parsed, dict) else None
        except (ValueError, TypeError):
            return None
    return None


class StageBaselineLoader(BaselineLoader):
    """Load baseline metadata and deltas from the implicit Snowflake stage.

    The stage is ``@{validation_database}.VALIDATION.BASELINES``.
    Metadata (procedure, params_hash, success, error, row_counts) and
    ``table_deltas`` (for DML validation) are loaded here.  The actual
    result-set row data is read directly from the stage at comparison
    time by ``materialize_baseline_transient_table``.
    """

    _STAGE_NAME = "BASELINES"

    def __init__(self, connector: ConnectorBase, database: str) -> None:
        self._connector = connector
        self._database = database
        self._fq_schema = f"{self._database}.VALIDATION"
        self._stage_ref = f"@{self._fq_schema}.{self._STAGE_NAME}"

    def load(self) -> list[BaselineData]:
        self._ensure_infrastructure()

        fmt_ref = f"{self._fq_schema}.JSON_FORMAT"
        try:
            rows = self._connector.execute_query(f"""\
SELECT
    $1:procedure::VARCHAR AS procedure_name,
    $1:params_hash::VARCHAR AS params_hash,
    $1:parameters AS params,
    $1:captured_at::VARCHAR AS captured_at,
    $1:success::BOOLEAN AS success,
    $1:error::VARCHAR AS error,
    $1:row_counts::ARRAY AS row_counts,
    $1:table_deltas AS table_deltas,
    $1:column_metadata AS column_metadata
FROM {self._stage_ref} (FILE_FORMAT => '{fmt_ref}')""")
        except Exception as e:
            print(f"Warning: Stage query failed: {e}")
            return []

        if not rows:
            return []

        baselines: list[BaselineData] = []
        for row in rows:
            try:
                params_raw = row.get("PARAMS") or row.get("params") or "{}"
                params = json.loads(params_raw) if isinstance(params_raw, str) else params_raw

                row_counts_raw = row.get("ROW_COUNTS") or row.get("row_counts")
                if row_counts_raw is not None:
                    try:
                        row_counts = json.loads(row_counts_raw) if isinstance(row_counts_raw, str) else list(row_counts_raw)
                    except (ValueError, TypeError):
                        row_counts = []
                else:
                    row_counts = []

                table_deltas = _parse_variant(
                    row.get("TABLE_DELTAS") or row.get("table_deltas")
                )

                col_meta_raw = _parse_variant(
                    row.get("COLUMN_METADATA") or row.get("column_metadata")
                )

                bl = BaselineData(
                    procedure=row.get("PROCEDURE_NAME") or row.get("procedure_name", ""),
                    parameters=params if isinstance(params, dict) else {},
                    captured_at=row.get("CAPTURED_AT") or row.get("captured_at", ""),
                    result_sets=[],
                    row_counts=row_counts,
                    success=row.get("SUCCESS", True) if row.get("SUCCESS") is not None else True,
                    error=row.get("ERROR") or row.get("error"),
                    stored_params_hash=row.get("PARAMS_HASH") or row.get("params_hash"),
                    table_deltas=table_deltas,
                    column_metadata=col_meta_raw,
                )
                baselines.append(bl)
            except Exception as exc:
                print(f"Warning: failed to parse baseline: {exc}")

        return baselines

    def _ensure_infrastructure(self) -> None:
        """Create database, schema, stage, and file format if they don't exist."""
        stmts = [
            f"CREATE DATABASE IF NOT EXISTS {self._database}",
            f"CREATE SCHEMA IF NOT EXISTS {self._fq_schema}",
            f"CREATE STAGE IF NOT EXISTS {self._fq_schema}.{self._STAGE_NAME} "
            f"COMMENT = 'Baseline JSON files'",
            f"CREATE FILE FORMAT IF NOT EXISTS {self._fq_schema}.JSON_FORMAT "
            f"TYPE = 'JSON'",
        ]
        for sql in stmts:
            try:
                self._connector.execute_statement(sql)
            except Exception:
                pass

    @property
    def source_label(self) -> str:
        return f"{self._stage_ref} (stage)"


# ---------------------------------------------------------------------------
# Factory function
# ---------------------------------------------------------------------------

def create_baseline_loader(
    connector: Optional[ConnectorBase],
    database: str,
    baseline_dir: Optional[Path],
    baseline_stage: Optional[str],
    project_root: Path,
) -> BaselineLoader:
    """Select the appropriate baseline loader based on configuration.

    When a Snowflake connector is available the implicit stage
    ``@{database}.VALIDATION.BASELINES`` is used.  Falls back to
    local disk for offline / testing scenarios.
    """
    if connector is not None:
        stage_loader = StageBaselineLoader(connector, database)
        baselines = stage_loader.load()
        if baselines:
            return _PreloadedLoader(baselines, stage_loader.source_label)

    if baseline_dir is None:
        baseline_dir = project_root / ".scai" / "baselines"
    return DiskBaselineLoader(baseline_dir)


class _PreloadedLoader(BaselineLoader):
    """Wraps already-loaded baselines to avoid loading twice."""

    def __init__(self, baselines: list[BaselineData], label: str) -> None:
        self._baselines = baselines
        self._label = label

    def load(self) -> list[BaselineData]:
        return self._baselines

    @property
    def source_label(self) -> str:
        return self._label
