"""
Capture runner -- orchestrates the full baseline-capture pipeline.

1. Resolve the source executor factory from the injected registry
2. Discover test YAML files from ``artifacts/**/test/*.yml``
3. For each test file / test case:
   a. Render the SQL template with dialect-specific literal formatting
   b. Execute on the source database
   c. Save baseline JSON to disk
4. Optionally upload baselines to a Snowflake stage
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional

from test_runner.common.baseline import write_baseline
from test_runner.common.config import ReplayConfiguration, TestRunnerConfig
from test_runner.common.database_dialect import DatabaseDialect
from test_runner.common.database_executor import DatabaseExecutor, DatabaseExecutorFactory
from test_runner.common.dependency_classifier import enrich_test_files
from test_runner.common.discovery import discover_test_files, filter_test_files
from test_runner.common.models import BaselineData, TestCaseStats, TestFile, ValidationSteps
from test_runner.common.template import (
    build_parameters,
    render_capture_sql,
    render_output_fetch_statements,
)
from test_runner.replay.capture_handler import capture_deltas
from test_runner.parallel.partitioner import partition_cases
from test_runner.parallel.executor import ParallelExecutor
from .stage_uploader import StageUploader

LOGGER = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Console output helpers
# ---------------------------------------------------------------------------

_SEPARATOR = "\u2550" * 60


def _print_header(dialect: str, baseline_dir: Path, stage: str = "") -> None:
    print(f"\n{_SEPARATOR}")
    print(f"  CAPTURING BASELINES FROM {dialect.upper()}")
    if stage:
        print(f"  Stage:  {stage}")
    print(f"  Output: {baseline_dir}")
    print(f"{_SEPARATOR}\n")


def _print_case_result(procedure: str, case_idx: int, result: BaselineData) -> None:
    icon = "\u2705" if result.success else "\u274c"
    if result.success:
        parts: list[str] = []
        if result.output_params:
            n = len(result.output_params)
            parts.append(f"{n} output param{'s' if n != 1 else ''}")
        if result.row_counts:
            if len(result.row_counts) == 1:
                parts.append(f"{result.row_counts[0]} rows")
            else:
                total = sum(result.row_counts)
                parts.append(f"{total} rows ({len(result.row_counts)} sets)")
        suffix = ", ".join(parts) if parts else "0 rows"
    else:
        suffix = f"ERROR: {result.error}"
    print(f"{icon} {procedure} (case {case_idx + 1}) -> {suffix}")


def _print_summary(stats: TestCaseStats) -> None:
    print(f"\n{_SEPARATOR}")
    print(f"  CAPTURED {stats.total} BASELINES")
    if stats.failed:
        print(f"  {stats.succeeded} succeeded, {stats.failed} failed")
    print(f"{_SEPARATOR}\n")


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def run_capture(
    config: TestRunnerConfig,
    factory_registry: dict[DatabaseDialect, DatabaseExecutorFactory],
    project_root: str | Path,
    baseline_dir: Optional[str | Path] = None,
    snowflake_connection: Optional[str] = None,
    workers: int = 1,
    where: str | None = None,
    include_dependencies: bool = False,
) -> TestCaseStats:
    """Execute the full capture pipeline.

    Baselines are always written to disk.  When a Snowflake connection is
    available they are also uploaded to the implicit internal stage
    ``@{validation_database}.VALIDATION.BASELINES`` and materialized
    into the ``BASELINE_ROWS`` table.

    Args:
        config: Loaded test runner configuration.
        factory_registry: Dialect -> factory mapping (injected).
        project_root: Path to the SCAI project root.
        baseline_dir: Override directory for baseline files.
        snowflake_connection: Snowflake connection name for stage upload.

    Returns:
        :class:`TestCaseStats` with totals and the list of written baseline paths.
    """
    project_root = Path(project_root)

    # -- resolve factory from registry ---------------------------------------
    factory = factory_registry.get(config.source_dialect)
    if factory is None:
        raise ValueError(
            f"No executor factory for dialect '{config.source_dialect.value}'."
        )

    # -- create executor via factory -----------------------------------------
    if not config.source_connection_raw:
        raise ValueError(
            "Capture requires a source connection. Configure a default in "
            "~/.snowflake/snowct/<dialect>.toml or pass --source-connection."
        )
    source_config = factory.build_config(config.source_connection_raw)
    factory.executor = factory.create_executor(source_config)

    factory.format_literal = factory.create_literal_formatter().format_literal

    # -- baseline directory --------------------------------------------------
    if baseline_dir is None:
        baseline_dir = project_root / ".scai" / "baselines"
    baseline_dir = Path(baseline_dir)

    capture_date = datetime.now(timezone.utc)

    # -- discover test files -------------------------------------------------
    test_files = discover_test_files(project_root, config.selectors)
    test_files = filter_test_files(test_files, project_root, where, include_dependencies)
    if not test_files:
        print("No test files found in artifacts/**/test/*.yml")
        return TestCaseStats()

    enrich_test_files(test_files, project_root)

    # -- stage uploader (uses provided stage or implicit default) --------------
    uploader: Optional[StageUploader] = None
    stage = f"@{config.validation_database}.VALIDATION.BASELINES"
    sf_factory = factory_registry.get(DatabaseDialect.SNOWFLAKE)
    if sf_factory is not None:
        target_raw = dict(config.target_connection_raw) if config.target_connection_raw else {}
        if snowflake_connection:
            target_raw["name"] = snowflake_connection
        if target_raw.get("name"):
            target_raw.setdefault("mode", "name")
            sf_config = sf_factory.build_config(target_raw)
            sf_executor = sf_factory.create_executor(sf_config)
            if sf_executor.connector is not None:
                uploader = StageUploader(
                    sf_executor.connector, stage, config.validation_database,
                )
            else:
                sf_executor.close()

    # -- run -----------------------------------------------------------------
    stats = TestCaseStats()

    _print_header(factory.dialect.value, baseline_dir, stage if uploader else "")

    try:
        if workers > 1:
            _run_parallel_capture(
                test_files=test_files,
                factory=factory,
                source_config=source_config,
                baseline_dir=baseline_dir,
                capture_date=capture_date,
                stats=stats,
                replay_cfg=config.replay_configuration,
                workers=workers,
            )
        else:
            for test_file in test_files:
                _capture_test_file(
                    test_file=test_file,
                    factory=factory,
                    baseline_dir=baseline_dir,
                    capture_date=capture_date,
                    stats=stats,
                    replay_cfg=config.replay_configuration,
                )
    finally:
        factory.executor.close()

    # -- stage upload -----------------------------------------------------------
    if uploader is not None:
        try:
            date_folder = capture_date.strftime("%Y%m%d")
            upload_root = baseline_dir / date_folder
            if upload_root.exists():
                count = uploader.upload_directory(upload_root, date_folder)
                print(f"\nUploaded {count} baselines to {stage}")
        finally:
            uploader.close()

    _print_summary(stats)
    return stats


# ---------------------------------------------------------------------------
# Internal
# ---------------------------------------------------------------------------


@dataclass
class _CaptureResult:
    """Result of capturing a single test case (for parallel aggregation)."""
    success: bool
    baseline_path: str


def _run_parallel_capture(
    test_files: list[TestFile],
    factory: DatabaseExecutorFactory,
    source_config: Any,
    baseline_dir: Path,
    capture_date: datetime,
    stats: TestCaseStats,
    replay_cfg: ReplayConfiguration | None,
    workers: int,
) -> None:
    """Partition test cases and run read-only captures in parallel.

    DML cases are executed sequentially using the existing
    ``_capture_test_file`` function.
    """
    dialect = factory.dialect.value
    format_literal = factory.format_literal

    read_only_cases, dml_cases = partition_cases(test_files)
    LOGGER.info(
        "[parallel] %d read-only cases, %d DML cases, %d workers",
        len(read_only_cases), len(dml_cases), workers,
    )

    # -- Read-only cases in parallel ------------------------------------------
    if read_only_cases:
        if source_config is None:
            raise RuntimeError(
                "Parallel capture requires a source factory and config."
            )

        def _worker_capture(
            thread_executor: DatabaseExecutor,
            test_file: TestFile,
            case_idx: int,
            test_case: list[Any],
        ) -> _CaptureResult:
            sql: str = render_capture_sql(test_file.source, test_case, dialect, format_literal)
            fetch_stmts: list[str] = render_output_fetch_statements(test_file.source, test_case, dialect)
            rendered_steps = ValidationSteps(
                run=test_file.source.run,
                pre_execute=test_file.source.pre_execute,
                output_parameters=test_file.source.output_parameters,
                output_tables=test_file.source.output_tables,
                output_cursors=test_file.source.output_cursors,
            )

            result = thread_executor.execute(
                sql, steps=rendered_steps, fetch_stmts=fetch_stmts or None,
            )

            parameters = build_parameters(test_file.parameter_names, test_case)
            baseline = BaselineData.from_capture(
                procedure=test_file.code_unit_name,
                parameters=parameters,
                result=result,
            )

            path = write_baseline(baseline, baseline_dir, case_idx, capture_date)
            _print_case_result(test_file.code_unit_name, case_idx, baseline)

            return _CaptureResult(success=result.success, baseline_path=str(path))

        par = ParallelExecutor(max_workers=workers)
        read_only_results = par.execute_with_factory(
            read_only_cases, factory, source_config, _worker_capture,
        )

        for capture_result in read_only_results:
            stats.total += 1
            if capture_result.success:
                stats.succeeded += 1
            else:
                stats.failed += 1
            stats.baseline_files.append(capture_result.baseline_path)

    # -- DML cases sequentially (grouped by test file) -------------------------
    # Group DML cases back by test file to use existing _capture_test_file.
    executed_tests_with_dml: set[str] = set()
    for test_file, _case_idx, _test_case in dml_cases:
        if test_file.file_path in executed_tests_with_dml:
            continue
        executed_tests_with_dml.add(test_file.file_path)
        _capture_test_file(
            test_file=test_file,
            factory=factory,
            baseline_dir=baseline_dir,
            capture_date=capture_date,
            stats=stats,
            replay_cfg=replay_cfg,
        )


def _capture_test_file(
    test_file: TestFile,
    factory: DatabaseExecutorFactory,
    baseline_dir: Path,
    capture_date: datetime,
    stats: TestCaseStats,
    replay_cfg: ReplayConfiguration | None = None,
) -> None:
    """Capture all test cases for a single test file."""
    executor = factory.executor
    dialect = factory.dialect.value

    outputs = test_file.source.output_parameters
    is_dml = (
        replay_cfg is not None
        and test_file.modifies_data is True
        and bool(test_file.affected_tables)
    )
    effective_replay_cfg = replay_cfg or ReplayConfiguration()

    for case_idx, test_case in enumerate(test_file.test_cases):
        sql = render_capture_sql(test_file.source, test_case, dialect, factory.format_literal)

        fetch_stmts = render_output_fetch_statements(
            test_file.source, test_case, dialect,
        )

        rendered_steps = ValidationSteps(
            run=test_file.source.run,
            pre_execute=test_file.source.pre_execute,
            output_parameters=outputs,
            output_tables=test_file.source.output_tables,
            output_cursors=test_file.source.output_cursors,
        )

        parameters = build_parameters(test_file.parameter_names, test_case)
        table_deltas: dict[str, dict] = {}
        column_metadata: dict[str, dict[str, list[str]]] = {}

        if is_dml:
            connection = getattr(executor, "connection", None)
            if connection is None and executor.connector is not None:
                connection = executor.connector.connection

            if connection is not None:
                def _exec():
                    return executor.execute(
                        sql, steps=rendered_steps,
                        fetch_stmts=fetch_stmts or None,
                    )

                try:
                    result, table_deltas, column_metadata = capture_deltas(
                        connection=connection,
                        affected_tables=test_file.affected_tables or [],
                        replay_cfg=effective_replay_cfg,
                        execute_fn=_exec,
                        executor_factory=factory,
                        code_unit_name=test_file.code_unit_name,
                    )
                except Exception as exc:
                    LOGGER.error(
                        "Delta capture failed for %s (case %d): %s — skipping case to preserve source integrity",
                        test_file.code_unit_name, case_idx, exc,
                    )
                    stats.total += 1
                    stats.failed += 1
                    continue
            else:
                LOGGER.warning(
                    "Cannot capture deltas for %s: no raw connection available",
                    test_file.code_unit_name,
                )
                result = executor.execute(
                    sql, steps=rendered_steps, fetch_stmts=fetch_stmts or None,
                )
        else:
            result = executor.execute(
                sql, steps=rendered_steps, fetch_stmts=fetch_stmts or None,
            )

        baseline = BaselineData.from_capture(
            procedure=test_file.code_unit_name,
            parameters=parameters,
            result=result,
        )
        if table_deltas:
            baseline.table_deltas = table_deltas
        if column_metadata:
            baseline.column_metadata = column_metadata

        path = write_baseline(baseline, baseline_dir, case_idx, capture_date)

        stats.total += 1
        if result.success:
            stats.succeeded += 1
        else:
            stats.failed += 1
        stats.baseline_files.append(str(path))

        _print_case_result(test_file.code_unit_name, case_idx, baseline)
