"""
DVT Model Runner — extends dbt's ModelRunner with cross-engine extraction.

Two-way dispatch:
1. Pushdown → delegates to stock dbt ModelRunner (all sources local)
2. Extraction → Sling → DuckDB cache → model SQL → Sling → target

ALL extraction models go through the persistent DuckDB cache (.dvt/cache.duckdb).
There is no "Sling Direct" sub-path. This ensures:
- Consistent DuckDB SQL dialect for all cross-engine models
- Incremental support via persistent cache
- Cache per source (shared across models)
"""

import logging
import time
import threading
from typing import Any, Dict, List, Optional, Tuple

from dbt.artifacts.schemas.results import RunStatus
from dbt.artifacts.schemas.run.v5.run import RunResult
from dbt.task.run import ModelRunner

from dvt.config.target_resolver import ExecutionPath, ResolvedModel

logger = logging.getLogger(__name__)


def _make_run_result(
    node: Any,
    status: RunStatus,
    message: str,
    execution_time: float,
    rows_affected: int = 0,
) -> RunResult:
    """Build a RunResult for extraction execution paths."""
    return RunResult(
        status=status,
        thread_id=threading.current_thread().name,
        execution_time=execution_time,
        timing=[],
        message=message,
        node=node,
        adapter_response={"_message": message, "rows_affected": rows_affected},
        failures=0 if status == RunStatus.Success else 1,
        batch_results=None,
    )


class DvtModelRunner(ModelRunner):
    """Extends ModelRunner with DVT extraction via DuckDB cache.

    Two-way dispatch:
    - Pushdown: super().execute() (stock dbt)
    - Extraction: Sling→DuckDB cache→Sling→target
    """

    _resolution: Optional[ResolvedModel] = None
    _profiles: Optional[dict] = None
    _cache = None  # DvtCache instance, shared across runners

    def set_dvt_context(
        self,
        resolution: ResolvedModel,
        profiles: dict,
        cache: Any = None,
    ) -> None:
        """Set DVT-specific context for this runner."""
        self._resolution = resolution
        self._profiles = profiles
        self._cache = cache

    def execute(self, model, manifest):
        """Two-way dispatch: pushdown vs extraction."""
        if self._resolution is None:
            return super().execute(model, manifest)

        path = self._resolution.execution_path

        if path == ExecutionPath.DEFAULT_PUSHDOWN:
            logger.info(f"DVT [{model.name}]: default_pushdown")
            return super().execute(model, manifest)

        elif path == ExecutionPath.NON_DEFAULT_PUSHDOWN:
            # Non-default target with all sources local to that target.
            # We can't use the default adapter (wrong engine). Route through
            # the extraction path: Sling extracts from non-default target →
            # DuckDB → Sling loads back to the same non-default target.
            logger.info(
                f"DVT [{model.name}]: non_default_pushdown via extraction "
                f"(target={self._resolution.target})"
            )
            return self._execute_extraction(model, manifest)

        elif path in (ExecutionPath.SLING_DIRECT, ExecutionPath.DUCKDB_COMPUTE):
            # Coerce view/ephemeral to table for extraction models (DVT001)
            materialization = model.get_materialization()
            if materialization in ("view", "ephemeral"):
                logger.warning(
                    f"DVT001: Model '{model.name}' is '{materialization}' but requires "
                    f"cross-engine extraction. Coercing to 'table'."
                )
                # Force table materialization
                model.config["materialized"] = "table"

            logger.info(f"DVT [{model.name}]: extraction via DuckDB cache")
            return self._execute_extraction(model, manifest)

        else:
            return super().execute(model, manifest)

    # ------------------------------------------------------------------
    # Unified Extraction: Sling → DuckDB cache → Sling → Target
    # ------------------------------------------------------------------

    def _execute_extraction(self, model, manifest) -> RunResult:
        """Unified extraction path for ALL cross-engine models.

        1. Extract each remote source → DuckDB cache (Sling)
        2. Rewrite compiled SQL (source refs → cache table names)
        3. Execute model SQL in DuckDB
        4. Store model result in DuckDB cache (for incremental {{ this }})
        5. Close DuckDB (release file lock)
        6. Sling loads result from DuckDB → target
        """
        from dvt.extraction.sling_client import SlingClient
        from dvt.federation.dvt_cache import DvtCache
        from dvt.config.source_connections import load_source_connections

        resolution = self._resolution
        start = time.time()

        try:
            client = SlingClient()
            cache = self._cache
            if cache is None:
                project_dir = getattr(self.config.args, "PROJECT_DIR", None) or "."
                cache = DvtCache(project_dir=project_dir)

            # Acquire file lock — DuckDB has exclusive file locking, and Sling
            # is an external process that also needs it. Serialize all extraction.
            DvtCache._file_lock.acquire()
            try:
                return self._execute_extraction_locked(
                    model, manifest, resolution, client, cache
                )
            finally:
                DvtCache._file_lock.release()

        except Exception as e:
            elapsed = time.time() - start
            error_detail = str(e)
            target_name = resolution.target if resolution else "unknown"
            target_type = ""
            try:
                tc = self._get_output_config(target_name)
                target_type = tc.get("type", "")
            except Exception:
                pass
            error_msg = (
                f"Model '{model.name}' (target: {target_name}, adapter: {target_type}): "
                f"{error_detail}\n"
                f"    Check logs/dvt.log for full details."
            )
            logger.error(error_msg)
            return _make_run_result(model, RunStatus.Error, error_msg, elapsed)

    def _execute_extraction_locked(
        self, model, manifest, resolution, client, cache
    ) -> RunResult:
        """Inner extraction logic — called while holding the file lock."""
        from dvt.config.source_connections import load_source_connections

        start = time.time()

        try:
            # Ensure cache file exists before any Sling operations
            cache.ensure_created()

            compiled_sql = getattr(model, "compiled_code", "") or ""
            if not compiled_sql:
                raise RuntimeError(f"Model {model.name} has no compiled SQL")

            # Collect ALL source dependencies
            depends_on = getattr(model, "depends_on", None)
            all_source_uids = []
            if depends_on:
                all_source_uids = [
                    n
                    for n in (getattr(depends_on, "nodes", []) or [])
                    if n.startswith("source.")
                ]

            project_dir = getattr(self.config.args, "PROJECT_DIR", None) or "."
            source_connections = load_source_connections(project_dir)

            # Check incremental state
            materialization = model.get_materialization()
            is_incremental = (
                materialization == "incremental" and cache.has_model_result(model.name)
            )

            # Resolve watermark if incremental
            watermark_value = None
            watermark_column = None
            if is_incremental:
                watermark_value, watermark_column = self._resolve_watermark_from_target(
                    model
                )

            # ----------------------------------------------------------
            # Step 1: Collect source metadata + rewrite SQL refs
            # ----------------------------------------------------------
            source_meta = []  # List of (source_uid, src_name, tbl_name, src_schema, src_identifier, connection, source_config, cache_table)
            ref_replacements = {}

            for source_uid in all_source_uids:
                source_node = manifest.sources.get(source_uid)
                if not source_node:
                    continue

                src_name = getattr(source_node, "source_name", "")
                tbl_name = getattr(source_node, "name", "")
                src_schema = getattr(source_node, "schema", "")
                src_identifier = getattr(source_node, "identifier", "") or tbl_name
                src_db = getattr(source_node, "database", "")

                connection = source_connections.get(src_name, resolution.target)
                source_config = self._get_output_config(connection)
                cache_table = cache.source_table_name(src_name, tbl_name)

                source_meta.append(
                    (
                        source_uid,
                        src_name,
                        tbl_name,
                        src_schema,
                        src_identifier,
                        connection,
                        source_config,
                        cache_table,
                    )
                )

                # Collect ref replacements
                candidates = []
                if src_db and src_schema:
                    candidates.append(f'"{src_db}"."{src_schema}"."{src_identifier}"')
                    candidates.append(f"{src_db}.{src_schema}.{src_identifier}")
                if src_schema:
                    candidates.append(f'"{src_schema}"."{src_identifier}"')
                    candidates.append(f"{src_schema}.{src_identifier}")

                for candidate in candidates:
                    if candidate in compiled_sql:
                        ref_replacements[candidate] = cache_table
                        break

            # Rewrite SQL (source refs → cache table names)
            rewritten_sql = compiled_sql
            for old_ref, new_ref in sorted(
                ref_replacements.items(), key=lambda x: len(x[0]), reverse=True
            ):
                rewritten_sql = rewritten_sql.replace(old_ref, new_ref)

            # ----------------------------------------------------------
            # Step 2: Run federation optimizer on rewritten SQL
            # ----------------------------------------------------------
            from dvt.federation.optimizer import optimize_extractions

            source_table_map = {
                sm[7]: sm[7] for sm in source_meta
            }  # cache_table → cache_table
            optimized = optimize_extractions(rewritten_sql, source_table_map)

            # ----------------------------------------------------------
            # Step 3: Extract each source → DuckDB cache (with optimization)
            # ----------------------------------------------------------
            for (
                source_uid,
                src_name,
                tbl_name,
                src_schema,
                src_identifier,
                connection,
                source_config,
                cache_table,
            ) in source_meta:
                # Build extraction query — use optimizer if available
                opt = optimized.get(cache_table)
                if opt and (opt.columns or opt.predicates or opt.limit is not None):
                    from dvt.federation.optimizer import ADAPTER_TO_SQLGLOT_DIALECT

                    source_type = source_config.get("type", "")
                    source_dialect = ADAPTER_TO_SQLGLOT_DIALECT.get(
                        source_type, "duckdb"
                    )
                    extraction_query = opt.build_query(
                        src_schema, src_identifier, source_dialect
                    )
                    logger.info(
                        f"DVT [{model.name}]: optimized → {source_dialect}: {extraction_query[:120]}"
                    )
                    extraction_query = opt.build_query(
                        src_schema, src_identifier, source_dialect
                    )
                    logger.info(
                        f"DVT [{model.name}]: optimized extraction for {cache_table} ({source_type}→{source_dialect}): {extraction_query[:100]}"
                    )
                else:
                    extraction_query = (
                        f"SELECT * FROM {src_schema}.{src_identifier}"
                        if src_schema
                        else f"SELECT * FROM {src_identifier}"
                    )

                # For incremental: add watermark filter
                if is_incremental and watermark_value is not None and watermark_column:
                    from dvt.extraction.watermark_formatter import format_watermark

                    source_type = source_config.get("type", "")
                    formatted = format_watermark(
                        watermark_value, source_type, "timestamp"
                    )
                    if "WHERE" in extraction_query.upper():
                        extraction_query += f" AND {watermark_column} > {formatted}"
                    else:
                        extraction_query += f" WHERE {watermark_column} > {formatted}"

                # Extract via Sling → DuckDB cache
                cache.close_and_release()
                client.extract_to_duckdb(
                    source_config=source_config,
                    duckdb_path=cache.db_path,
                    source_query=extraction_query,
                    duckdb_table=cache_table,
                )
                cache.reopen()

            # Also extract remote refs (no optimization for refs)
            for ref_uid in resolution.remote_refs or []:
                ref_node = manifest.nodes.get(ref_uid)
                if not ref_node:
                    continue
                ref_target = (
                    getattr(ref_node.config, "target", None) or self.config.target_name
                )
                ref_config = self._get_output_config(ref_target)
                ref_schema = getattr(ref_node, "schema", "")
                ref_name = getattr(ref_node, "name", "")
                query = (
                    f"SELECT * FROM {ref_schema}.{ref_name}"
                    if ref_schema
                    else f"SELECT * FROM {ref_name}"
                )

                cache.close_and_release()
                client.extract_to_duckdb(
                    source_config=ref_config,
                    duckdb_path=cache.db_path,
                    source_query=query,
                    duckdb_table=ref_name,
                )
                cache.reopen()

            # Use the already-rewritten SQL for DuckDB execution
            compiled_sql = rewritten_sql

            # Always rewrite {{ this }} if the compiled SQL references the model itself.
            # This handles both cases:
            # - is_incremental=true (cache has previous result)
            # - is_incremental=false but compiled SQL still has {{ this }} from stale partial_parse
            # The compiled SQL references the target table (e.g., "devdb"."public"."model_name").
            # In DuckDB cache, it's __model__model_name.
            if True:  # Always attempt {{ this }} rewriting
                model_cache_table = cache.model_table_name(model.name)
                target_table_fqn = self._model_table_name(model)
                # Try various formats dbt might have generated for {{ this }}
                this_candidates = []
                if model.database and model.schema:
                    this_candidates.append(
                        f'"{model.database}"."{model.schema}"."{model.name}"'
                    )
                    this_candidates.append(
                        f"{model.database}.{model.schema}.{model.name}"
                    )
                if model.schema:
                    this_candidates.append(f'"{model.schema}"."{model.name}"')
                    this_candidates.append(f"{model.schema}.{model.name}")
                # Only use qualified names (schema.name) for {{ this }} replacement.
                # The bare model name can match inside other identifiers
                # (e.g., "f_country" inside "cbs_f_country") causing corruption.
                for candidate in this_candidates:
                    if candidate in compiled_sql:
                        compiled_sql = compiled_sql.replace(
                            candidate, model_cache_table
                        )
                        break

            # ----------------------------------------------------------
            # Step 3: Execute model SQL in DuckDB + cache the result
            # ----------------------------------------------------------
            # If the model references {{ this }} (rewritten to __model__xxx)
            # but the cache doesn't have the model result table yet,
            # this is a stale is_incremental() from dbt's partial_parse.
            # Treat as first run: strip the incremental WHERE clause.
            model_cache_table = cache.model_table_name(model.name)
            if model_cache_table in compiled_sql and not cache.has_model_result(
                model.name
            ):
                logger.info(
                    f"DVT [{model.name}]: stale is_incremental() detected — "
                    f"cache has no {model_cache_table}, stripping WHERE clause"
                )
                # Remove everything from WHERE to end if it references the cache table
                import re

                # Strip WHERE ... that references the model cache table
                compiled_sql = re.sub(
                    r"\bWHERE\b.*" + re.escape(model_cache_table) + r".*$",
                    "",
                    compiled_sql,
                    flags=re.IGNORECASE | re.DOTALL,
                )

            # For append strategy, accumulate rows in cache instead of replacing
            incremental_strategy = ""
            if is_incremental:
                incremental_strategy = model.config.get(
                    "incremental_strategy", ""
                )
            use_append = is_incremental and incremental_strategy == "append"

            logger.info(f"DVT [{model.name}]: DuckDB SQL: {compiled_sql[:200]}...")
            row_count = cache.save_model_result(
                model.name, compiled_sql, append=use_append
            )

            # ----------------------------------------------------------
            # Step 4: Close DuckDB, Sling loads result → target
            # ----------------------------------------------------------
            result_table = cache.model_table_name(model.name)
            db_path = cache.db_path
            cache.close_and_release()

            target_config = self._get_output_config(resolution.target)
            target_table = self._model_table_name(model)

            # Determine Sling mode for target
            update_key = None
            if is_incremental:
                unique_key = model.config.get("unique_key")
                incremental_strategy = model.config.get(
                    "incremental_strategy", ""
                )

                if incremental_strategy == "append" or not unique_key:
                    # Append strategy: DuckDB cache already has the full result
                    # (incremental WHERE already applied during SQL execution).
                    # Sling's incremental/backfill modes require primary_key
                    # which append models don't have. Use full-refresh to
                    # replace the target with the complete cached result.
                    sling_mode = "full-refresh"
                    primary_key = None
                else:
                    sling_mode = "incremental"
                    primary_key = (
                        [unique_key]
                        if isinstance(unique_key, str)
                        else (unique_key or [])
                    )
            else:
                # full-refresh drops the table, which fails when dependent
                # views exist. Use truncate instead — it clears and reloads
                # without dropping, preserving dependent views.
                sling_mode = "truncate"
                primary_key = None

            client.load_from_duckdb(
                duckdb_path=db_path,
                target_config=target_config,
                source_query=f"SELECT * FROM {result_table}",
                target_table=target_table,
                mode=sling_mode,
                primary_key=primary_key,
                update_key=update_key,
            )

            # Reopen cache for potential next model
            cache.reopen()

            elapsed = time.time() - start
            incr_label = " (incremental)" if is_incremental else ""
            msg = f"SELECT {row_count}{incr_label}"
            return _make_run_result(model, RunStatus.Success, msg, elapsed, row_count)

        except Exception as e:
            elapsed = time.time() - start
            error_detail = str(e)
            target_name = resolution.target if resolution else "unknown"
            target_type = ""
            try:
                tc = self._get_output_config(target_name)
                target_type = tc.get("type", "")
            except Exception:
                pass
            error_msg = (
                f"Model '{model.name}' (target: {target_name}, adapter: {target_type}): "
                f"{error_detail}\n"
                f"    Check logs/dvt.log for full details."
            )
            logger.error(error_msg)
            return _make_run_result(model, RunStatus.Error, error_msg, elapsed)

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _resolve_watermark_from_target(
        self, model
    ) -> Tuple[Optional[Any], Optional[str]]:
        """Query the TARGET for the incremental watermark value.

        Scans the compiled SQL for common update_key column names.
        Returns (watermark_value, column_name) or (None, None).
        """
        compiled = getattr(model, "compiled_code", "") or ""

        # Find update_key column from common patterns
        update_key = None
        for candidate in [
            "updated_at",
            "modified_at",
            "_updated_at",
            "last_modified",
            "update_date",
            "modified_date",
            "created_at",
            "loaded_at",
        ]:
            if candidate in compiled.lower():
                update_key = candidate
                break

        if not update_key:
            return None, None

        try:
            target_table = self._model_table_name(model)
            sql = f"SELECT MAX({update_key}) FROM {target_table}"
            _, result = self.adapter.execute(sql, fetch=True)
            if result and len(result) > 0 and result[0][0] is not None:
                return result[0][0], update_key
        except Exception:
            pass

        return None, None

    def _model_table_name(self, model) -> str:
        """Get the fully-qualified target table name for a model.

        For non-default targets, uses the schema from the target's profiles.yml
        config instead of the model's compiled schema (which comes from the
        default adapter).
        """
        schema = model.schema

        # For any model targeting a non-default target: use the target's schema
        # from profiles.yml. The default adapter compiles schema as 'public' (PG),
        # but MSSQL uses 'dbo', Oracle uses 'SYSTEM', etc.
        if self._resolution and self._resolution.target != self.config.target_name:
            target_config = self._get_output_config(self._resolution.target)
            target_schema = target_config.get(
                "schema", target_config.get("database", "")
            )
            if target_schema:
                schema = target_schema

        if schema:
            return f"{schema}.{model.name}"
        return model.name

    def _get_output_config(self, output_name: str) -> dict:
        """Get a profiles.yml output config by name."""
        if not self._profiles:
            raise RuntimeError("DVT: profiles not set on runner")

        for profile_name, profile_data in self._profiles.items():
            if not isinstance(profile_data, dict):
                continue
            outputs = profile_data.get("outputs", {})
            if output_name in outputs:
                config = outputs[output_name]
                if isinstance(config, dict):
                    return {**config, "type": config.get("type", "")}

        raise RuntimeError(
            f"DVT101: Connection '{output_name}' not found in profiles.yml"
        )
