"""
Deploy Snowpark validator to Snowflake.

Packages the self-contained ``snowpark/`` module into a zip, uploads it to a
Snowflake stage, and creates the ``VALIDATE_BATCH`` / ``VALIDATE_SINGLE``
stored procedures.
"""

from __future__ import annotations

import shutil
import tempfile
from pathlib import Path
from typing import Any

from snowflake.snowflake_data_validation.connector.connector_base import ConnectorBase

from .results_writer import create_validation_schema


# ---------------------------------------------------------------------------
# SQL templates for Snowpark stored procedures
# ---------------------------------------------------------------------------

_CREATE_VALIDATE_BATCH_SQL = """\
CREATE OR REPLACE PROCEDURE {database}.VALIDATION.VALIDATE_BATCH(
    BASELINE_STAGE VARCHAR,
    PATTERN VARCHAR,
    DATABASE_NAME VARCHAR,
    VALIDATION_SCHEMA VARCHAR DEFAULT 'VALIDATION',
    PROC_PREFIX VARCHAR DEFAULT ''
)
RETURNS VARIANT
LANGUAGE PYTHON
RUNTIME_VERSION = '3.10'
PACKAGES = ('snowflake-snowpark-python')
IMPORTS = ('{stage}/snowpark_validator.zip')
HANDLER = 'run_batch'
EXECUTE AS CALLER
AS $$
import sys
sys.path.insert(0, sys._xoptions.get("snowflake_import_directory", "/tmp"))

from snowpark_validator.main import validate_batch

def run_batch(session, baseline_stage, pattern, database_name, validation_schema, proc_prefix):
    return validate_batch(
        session, baseline_stage, pattern, database_name,
        validation_schema, proc_prefix,
        results_table=f"{{validation_schema}}.RESULTS"
    )
$$"""

_CREATE_VALIDATE_SINGLE_SQL = """\
CREATE OR REPLACE PROCEDURE {database}.VALIDATION.VALIDATE_SINGLE(
    PROCEDURE_NAME VARCHAR,
    PARAMS VARIANT,
    BASELINE_STAGE VARCHAR,
    DATABASE_NAME VARCHAR,
    VALIDATION_SCHEMA VARCHAR DEFAULT 'VALIDATION',
    PROC_PREFIX VARCHAR DEFAULT ''
)
RETURNS VARIANT
LANGUAGE PYTHON
RUNTIME_VERSION = '3.10'
PACKAGES = ('snowflake-snowpark-python')
IMPORTS = ('{stage}/snowpark_validator.zip')
HANDLER = 'run_single'
EXECUTE AS CALLER
AS $$
import sys
sys.path.insert(0, sys._xoptions.get("snowflake_import_directory", "/tmp"))

from snowpark_validator.main import validate_single

def run_single(session, procedure_name, params, baseline_stage, database_name, validation_schema, proc_prefix):
    return validate_single(
        session, procedure_name, params, baseline_stage,
        database_name, validation_schema, proc_prefix
    )
$$"""


# ---------------------------------------------------------------------------
# Packaging
# ---------------------------------------------------------------------------

def _get_snowpark_src_dir() -> Path:
    return Path(__file__).parent / "snowpark"


def create_package_zip(output_path: Path) -> Path:
    """Package the ``snowpark/`` module into ``snowpark_validator.zip``."""
    src_dir = _get_snowpark_src_dir()
    if not src_dir.exists():
        raise FileNotFoundError(f"Snowpark source directory not found: {src_dir}")

    with tempfile.TemporaryDirectory() as tmp:
        tmp_path = Path(tmp)
        package_dir = tmp_path / "snowpark_validator"
        package_dir.mkdir()

        for py_file in src_dir.glob("*.py"):
            shutil.copy(py_file, package_dir / py_file.name)

        zip_base = output_path.parent / output_path.stem
        shutil.make_archive(str(zip_base), "zip", tmp_path)

    return output_path


# ---------------------------------------------------------------------------
# Upload & deploy
# ---------------------------------------------------------------------------

def upload_package(connector: ConnectorBase, zip_path: Path, stage: str) -> None:
    """PUT the validator zip to the Snowflake stage."""
    stage = stage.rstrip("/")
    put_sql = f"PUT 'file://{zip_path}' '{stage}/' AUTO_COMPRESS=FALSE OVERWRITE=TRUE"
    connector.execute_statement(put_sql)


def create_snowpark_procedures(connector: ConnectorBase, stage: str, database: str) -> None:
    """Create ``VALIDATE_BATCH`` and ``VALIDATE_SINGLE`` stored procedures."""
    stage = stage.rstrip("/")
    connector.execute_statement(_CREATE_VALIDATE_BATCH_SQL.format(stage=stage, database=database))
    connector.execute_statement(_CREATE_VALIDATE_SINGLE_SQL.format(stage=stage, database=database))


# ---------------------------------------------------------------------------
# Full deploy (public)
# ---------------------------------------------------------------------------

def deploy_snowpark_validator(
    connector: ConnectorBase,
    database: str,
    stage: str | None = None,
) -> None:
    """Full deploy: schema + table + views + zip upload + stored procedures."""
    fq = f"{database}.VALIDATION"
    if stage is None:
        stage = f"@{fq}.CODE_STAGE"

    # 1. Create schema, table, views
    create_validation_schema(connector, database)

    # 2. Create internal stage
    stage_name = stage.lstrip("@")
    connector.execute_statement(
        f"CREATE STAGE IF NOT EXISTS {stage_name} "
        f"COMMENT = 'Snowpark validator code package'"
    )

    # 3. Package and upload the zip
    with tempfile.TemporaryDirectory() as tmp:
        zip_path = Path(tmp) / "snowpark_validator.zip"
        create_package_zip(zip_path)
        upload_package(connector, zip_path, stage)

    # 4. Create stored procedures
    create_snowpark_procedures(connector, stage, database)

    print(f"Deployed Snowpark validator to {stage}")
    print(f"  {fq}.VALIDATE_BATCH  -- batch validation")
    print(f"  {fq}.VALIDATE_SINGLE -- single procedure validation")
