"""
Upload baseline files to a Snowflake stage.

Uses ``ConnectorSnowflake`` from ``snowflake-data-validation`` to PUT local
baseline JSON files into an internal stage.
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any

from snowflake.snowflake_data_validation.connector.connector_base import ConnectorBase

from test_runner.validate.results_writer import create_validation_schema

LOGGER = logging.getLogger(__name__)


class StageUploader:
    """Uploads files to a Snowflake internal stage.

    Usage::

        uploader = StageUploader(connector, "@MYDB.MYSCHEMA.BASELINES")
        uploader.upload_file(Path("baselines/case_000.json"))
        uploader.close()
    """

    def __init__(self, connector: ConnectorBase, stage: str, database: str = "") -> None:
        self._connector = connector
        self._stage = stage.rstrip("/")
        self._database = database

    def close(self) -> None:
        self._connector.close()

    def __enter__(self) -> StageUploader:
        return self

    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        self.close()

    # -- upload --------------------------------------------------------------

    def upload_file(self, local_path: Path, stage_subpath: str = "") -> None:
        """PUT a single file to the stage."""
        target = self._stage
        if stage_subpath:
            target = f"{target}/{stage_subpath.strip('/')}"

        put_sql = f"PUT 'file://{local_path}' '{target}/' AUTO_COMPRESS=FALSE OVERWRITE=TRUE"
        self._connector.execute_statement(put_sql)

    def upload_directory(self, local_dir: Path, stage_subpath: str = "") -> int:
        """PUT all ``*.json`` files under *local_dir* to the stage.

        Ensures the validation schema exists first, preserves the relative
        directory structure, and returns the number of files uploaded.
        """
        if self._database:
            try:
                create_validation_schema(self._connector, self._database)
            except Exception:
                pass

        count = 0
        for json_file in sorted(local_dir.rglob("*.json")):
            relative = json_file.relative_to(local_dir)
            sub = "/".join([stage_subpath, str(relative.parent)]) if stage_subpath else str(relative.parent)
            self.upload_file(json_file, sub)
            count += 1
        return count
