"""Generic Dimensional Model Loader with parallel support.

Generalised from special_projects/load_enertia_dm_masked.py so any ERP's
table specs can be loaded into a target DW schema without custom code.

Supports:
- JSON spec files defining source→target column mappings
- Optional DataShield masking via ScrambleEngine
- Parallel dimension loading via ThreadPoolExecutor
- Batch inserts with configurable chunk size
"""

from __future__ import annotations

import json
import logging
import os
import re as _re
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Spec dataclasses
# ---------------------------------------------------------------------------

@dataclass
class ColumnSpec:
    """Mapping of one target column."""
    alias: str
    kind: str  # num, str, ts, bool, bool_inv, date_key, const, null_num, null_str, null_ts
    source_names: Optional[List[str]] = None
    const_value: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        d: Dict[str, Any] = {"alias": self.alias, "kind": self.kind}
        if self.source_names:
            d["source_names"] = self.source_names
        if self.const_value is not None:
            d["const_value"] = self.const_value
        return d

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "ColumnSpec":
        return cls(
            alias=d["alias"],
            kind=d["kind"],
            source_names=d.get("source_names"),
            const_value=d.get("const_value"),
        )

    @classmethod
    def from_tuple(cls, t: tuple) -> "ColumnSpec":
        """Parse legacy (alias, kind, [names], value?) tuple format."""
        alias, kind = t[0], t[1]
        names = t[2] if len(t) > 2 else None
        value = t[3] if len(t) > 3 else None
        return cls(alias=alias, kind=kind, source_names=names, const_value=value)


@dataclass
class TableSpec:
    """Full mapping for one target table."""
    target_table: str
    source_tables: List[str]
    is_transaction: bool = False
    columns: List[ColumnSpec] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "target_table": self.target_table,
            "source_tables": self.source_tables,
            "is_transaction": self.is_transaction,
            "columns": [c.to_dict() for c in self.columns],
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "TableSpec":
        return cls(
            target_table=d["target_table"],
            source_tables=d.get("source_tables", []),
            is_transaction=d.get("is_transaction", False),
            columns=[ColumnSpec.from_dict(c) for c in d.get("columns", [])],
        )


class DMSpecLoader:
    """Load TableSpec definitions from JSON files."""

    @staticmethod
    def from_json(path: str | Path) -> List[TableSpec]:
        """Parse a JSON spec file into TableSpec objects.

        The JSON should be an array of objects matching TableSpec.to_dict() format.
        """
        text = Path(path).read_text(encoding="utf-8")
        data = json.loads(text)
        if not isinstance(data, list):
            data = [data]
        return [TableSpec.from_dict(d) for d in data]

    @staticmethod
    def to_json(specs: List[TableSpec], path: Optional[str | Path] = None) -> str:
        """Serialize specs to JSON. Optionally write to file."""
        text = json.dumps([s.to_dict() for s in specs], indent=2)
        if path:
            Path(path).parent.mkdir(parents=True, exist_ok=True)
            Path(path).write_text(text, encoding="utf-8")
        return text


# ---------------------------------------------------------------------------
# SQL expression builder (generalised from load_enertia_dm_masked.py)
# ---------------------------------------------------------------------------

def _pick(available: set, candidates: Optional[List[str]]) -> Optional[str]:
    """Find first matching column from candidates in available columns."""
    if not candidates:
        return None
    for c in candidates:
        if c in available:
            return c
        if c.upper() in available:
            return c.upper()
        if c.lower() in available:
            return c.lower()
    return None


def _expr(colset: set, spec: ColumnSpec) -> str:
    """Build a SQL SELECT expression from a ColumnSpec."""
    alias = spec.alias
    kind = spec.kind

    if kind == "const":
        return f"{spec.const_value} AS {alias}"

    if kind.startswith("null_"):
        type_map = {"null_num": "NUMBER", "null_str": "STRING", "null_ts": "TIMESTAMP_NTZ"}
        return f"NULL::{type_map.get(kind, 'STRING')} AS {alias}"

    c = _pick(colset, spec.source_names)
    if not c:
        # Fallback to NULL of appropriate type
        fallback = {"num": "NUMBER", "str": "STRING", "ts": "TIMESTAMP_NTZ",
                     "bool": "BOOLEAN", "bool_inv": "BOOLEAN", "date_key": "NUMBER"}
        return f"NULL::{fallback.get(kind, 'STRING')} AS {alias}"

    if kind == "num":
        return f'TRY_TO_NUMBER(TO_VARCHAR("{c}")) AS {alias}'
    elif kind == "str":
        return f'TO_VARCHAR("{c}") AS {alias}'
    elif kind == "ts":
        # Direct cast handles TIMESTAMP_NTZ (identity), DATE (upcast), and
        # VARCHAR (implicit parse).  Avoids TRY_TO_TIMESTAMP_NTZ which
        # Snowflake converts to TRY_CAST and then rejects when source is
        # already TIMESTAMP_NTZ.
        return f'"{c}"::TIMESTAMP_NTZ AS {alias}'
    elif kind == "bool":
        return f'IFF(TRY_TO_NUMBER(TO_VARCHAR("{c}")) = 1, TRUE, FALSE) AS {alias}'
    elif kind == "bool_inv":
        return f'IFF(TRY_TO_NUMBER(TO_VARCHAR("{c}")) = 1, FALSE, TRUE) AS {alias}'
    elif kind == "date_key":
        # TRY_TO_DATE handles TIMESTAMP_NTZ→DATE (different types, no
        # TRY_CAST same-type error) and VARCHAR→DATE safely.
        return f"TRY_TO_NUMBER(TO_CHAR(TRY_TO_DATE(\"{c}\"::VARCHAR), 'YYYYMMDD'), 'YYYYMMDD') AS {alias}"
    else:
        return f'TO_VARCHAR("{c}") AS {alias}'


# Common date column names used in Enertia/ERP WHERE filters, ordered by
# preference.  When a static WHERE clause references columns that don't
# exist in the source table, _safe_where builds a COALESCE using only the
# columns that *do* exist.
_DATE_COL_CANDIDATES = [
    "TXNDATE", "TXNSVCDATE", "EFFECTIVEDATE", "GLDATE",
    "ALLOCDATE", "PROCESSDATE", "CREATEDATE",
]


def _safe_where(where_clause: str, colset: set) -> str:
    """Build a WHERE clause that only references columns present in *colset*.

    If ``where_clause`` contains a COALESCE over column names that may not
    exist in the source table, Snowflake raises a compile-time error.  This
    helper rebuilds the filter using only the date columns actually present.

    Falls back to returning the original clause unchanged when the pattern
    is not recognised (non-COALESCE filters).
    """
    # Only rewrite COALESCE-style date filters
    if "COALESCE(" not in where_clause.upper():
        return where_clause

    # Extract the comparison value (e.g. '2025-01-01'::TIMESTAMP_NTZ)
    m = _re.search(r">= *(.+)$", where_clause.strip(), _re.IGNORECASE)
    if not m:
        return where_clause
    threshold = m.group(1).strip()

    # Find which date columns actually exist in this source table
    upper_colset = {c.upper() for c in colset}
    present = [c for c in _DATE_COL_CANDIDATES if c in upper_colset]

    if not present:
        # No date columns at all — skip the WHERE entirely
        return ""

    if len(present) == 1:
        # Use TRY_TO_DATE via ::VARCHAR to avoid TRY_CAST same-type error
        # when the column is already TIMESTAMP_NTZ.
        return f'TRY_TO_DATE("{present[0]}"::VARCHAR) >= {threshold}'

    parts = ", ".join(f'TRY_TO_DATE("{c}"::VARCHAR)' for c in present)
    return f"COALESCE({parts}, {threshold}) >= {threshold}"


# ---------------------------------------------------------------------------
# Generic DM Loader
# ---------------------------------------------------------------------------

class GenericDMLoader:
    """Load dimensional model tables from source to destination.

    Supports optional DataShield masking and parallel dimension loading.
    When ``sf_source_conn`` or ``sf_dest_conn`` are provided, the loader
    reuses those connections instead of creating new ones (avoids SSO prompts).
    """

    def __init__(
        self,
        source_db: str,
        source_schema: str,
        dest_db: str,
        dest_schema: str,
        source_connection: Optional[str] = None,
        dest_connection: Optional[str] = None,
        sf_source_conn=None,
        sf_dest_conn=None,
    ):
        self.source_db = source_db
        self.source_schema = source_schema
        self.dest_db = dest_db
        self.dest_schema = dest_schema
        self.source_connection = source_connection
        self.dest_connection = dest_connection
        self._shared_source_conn = sf_source_conn
        self._shared_dest_conn = sf_dest_conn

    def _get_connector(self):
        try:
            import snowflake.connector
            return snowflake.connector
        except ImportError as exc:
            raise RuntimeError("snowflake-connector-python required") from exc

    def _connect_from_cli(self, connection_name: str, database: str, schema: str):
        """Connect to Snowflake using a named CLI connection from config.toml."""
        from .snowflake_utils import get_snowflake_connection
        connector = self._get_connector()
        params = get_snowflake_connection(connection_name)
        if not params or not params.get("account"):
            raise RuntimeError(f"Snowflake CLI connection not found or missing account: {connection_name}")
        return connector.connect(
            account=params.get("account", ""),
            user=params.get("user", ""),
            password=params.get("password", ""),
            role=params.get("role", ""),
            warehouse=params.get("warehouse", ""),
            authenticator=params.get("authenticator", ""),
            database=database,
            schema=schema,
            client_store_temporary_credential=True,
        )

    def _connect_source(self):
        if self._shared_source_conn is not None:
            return self._shared_source_conn
        if self.source_connection:
            return self._connect_from_cli(self.source_connection, self.source_db, self.source_schema)
        connector = self._get_connector()
        return connector.connect(
            account=os.getenv("SNOWFLAKE_ACCOUNT", ""),
            user=os.getenv("SNOWFLAKE_USER", ""),
            password=os.getenv("SNOWFLAKE_PASSWORD", ""),
            role=os.getenv("SNOWFLAKE_ROLE", ""),
            warehouse=os.getenv("SNOWFLAKE_WAREHOUSE", ""),
            database=self.source_db,
            schema=self.source_schema,
            client_store_temporary_credential=True,
        )

    def _connect_dest(self):
        if self._shared_dest_conn is not None:
            return self._shared_dest_conn
        if self.dest_connection:
            return self._connect_from_cli(self.dest_connection, self.dest_db, self.dest_schema)
        connector = self._get_connector()
        return connector.connect(
            account=os.getenv("SNOWFLAKE_DEST_ACCOUNT", os.getenv("SNOWFLAKE_ACCOUNT", "")),
            user=os.getenv("SNOWFLAKE_DEST_USER", os.getenv("SNOWFLAKE_USER", "")),
            password=os.getenv("SNOWFLAKE_DEST_PASSWORD", os.getenv("SNOWFLAKE_PASSWORD", "")),
            role=os.getenv("SNOWFLAKE_DEST_ROLE", os.getenv("SNOWFLAKE_ROLE", "")),
            warehouse=os.getenv("SNOWFLAKE_DEST_WAREHOUSE", os.getenv("SNOWFLAKE_WAREHOUSE", "")),
            database=self.dest_db,
            schema=self.dest_schema,
            client_store_temporary_credential=True,
        )

    def load(
        self,
        specs: List[TableSpec],
        dim_limit: int = 200_000,
        txn_limit: int = 1_000_000,
        batch_size: int = 10_000,
        truncate: bool = True,
        parallel: bool = False,
        max_workers: int = 4,
        txn_where: str = "",
    ) -> Dict[str, Any]:
        """Load all specs from source to destination.

        Args:
            specs: Table specifications to load.
            dim_limit: Row limit for dimension tables.
            txn_limit: Row limit for transaction/fact tables.
            batch_size: Insert batch size.
            truncate: Truncate target before loading.
            parallel: Load dimensions concurrently.
            max_workers: Thread pool size for parallel mode.
            txn_where: Optional WHERE clause for fact/transaction tables
                (e.g. ``"TXNDATE >= '2025-01-01'"``).  Applied only to
                transaction specs.

        Returns:
            Dict with per-table results and total duration.
        """
        t0 = time.time()
        dim_specs = [s for s in specs if not s.is_transaction]
        fact_specs = [s for s in specs if s.is_transaction]
        results: Dict[str, Dict[str, Any]] = {}

        if parallel and len(dim_specs) > 1:
            # Load dimensions in parallel
            with ThreadPoolExecutor(max_workers=max_workers) as pool:
                futures = {
                    pool.submit(
                        self._load_table, spec, dim_limit, batch_size, truncate
                    ): spec.target_table
                    for spec in dim_specs
                }
                for future in as_completed(futures):
                    name = futures[future]
                    try:
                        results[name] = future.result()
                    except Exception as exc:
                        results[name] = {"status": "error", "error": str(exc), "rows": 0}
        else:
            for spec in dim_specs:
                results[spec.target_table] = self._load_table(
                    spec, dim_limit, batch_size, truncate
                )

        # Facts always after dims (may reference dim keys)
        if parallel and len(fact_specs) > 1:
            with ThreadPoolExecutor(max_workers=min(max_workers, 3)) as pool:
                futures = {
                    pool.submit(
                        self._load_table, spec, txn_limit, batch_size, truncate, txn_where
                    ): spec.target_table
                    for spec in fact_specs
                }
                for future in as_completed(futures):
                    name = futures[future]
                    try:
                        results[name] = future.result()
                    except Exception as exc:
                        results[name] = {"status": "error", "error": str(exc), "rows": 0}
        else:
            for spec in fact_specs:
                results[spec.target_table] = self._load_table(
                    spec, txn_limit, batch_size, truncate, txn_where
                )

        total_rows = sum(r.get("rows", 0) for r in results.values())
        duration = round(time.time() - t0, 2)

        return {
            "status": "completed",
            "tables_loaded": len(results),
            "total_rows": total_rows,
            "duration_seconds": duration,
            "results": results,
        }

    def _load_table(
        self,
        spec: TableSpec,
        row_limit: int,
        batch_size: int,
        truncate: bool,
        where_clause: str = "",
    ) -> Dict[str, Any]:
        """Load a single table spec.

        When shared connections are provided, they are reused (no new SSO).
        Otherwise falls back to creating per-table connections.
        """
        t0 = time.time()
        target = spec.target_table
        logger.info("DM Load: starting %s ...", target)
        # Determine if we own the connections (and should close them) or are reusing shared ones
        own_src = self._shared_source_conn is None
        own_dst = self._shared_dest_conn is None

        try:
            src_conn = self._connect_source()
            dst_conn = self._connect_dest()
            try:
                src_cur = src_conn.cursor()
                dst_cur = dst_conn.cursor()

                # Fully-qualified source prefix (avoids session context issues)
                src_prefix = f'"{self.source_db}"."{self.source_schema}".' if self.source_db and self.source_schema else ""

                # Find available source table
                source = None
                for candidate in spec.source_tables:
                    try:
                        src_cur.execute(f'SELECT 1 FROM {src_prefix}"{candidate}" LIMIT 1')
                        source = candidate
                        break
                    except Exception:
                        continue

                if not source:
                    logger.info("DM Load: %s -- no source table found (tried %s)", target, spec.source_tables)
                    return {"status": "skipped", "reason": "no source table found", "rows": 0}

                logger.info("DM Load: %s <- %s (source found)", target, source)

                # Get source columns
                src_cur.execute(f'SELECT * FROM {src_prefix}"{source}" LIMIT 0')
                colset = {desc[0] for desc in src_cur.description}

                # Build expressions
                exprs = [_expr(colset, col) for col in spec.columns]
                base_sql = f'SELECT {", ".join(exprs)} FROM {src_prefix}"{source}"'
                if where_clause:
                    safe_where = _safe_where(where_clause, colset)
                    if safe_where:
                        base_sql = f'{base_sql} WHERE {safe_where}'
                    else:
                        logger.info("DM Load: %s -- WHERE clause skipped (no matching date columns in %s)", target, source)
                select_sql = f'{base_sql} LIMIT {row_limit}'

                # Create target if needed and optionally truncate.
                # DDL runs on DEST cursor so cannot reference cross-account source tables.
                # Derive column types from ColumnSpec.kind instead.
                if truncate:
                    kind_to_type = {
                        "num": "NUMBER", "str": "VARCHAR", "ts": "TIMESTAMP_NTZ",
                        "bool": "BOOLEAN", "bool_inv": "BOOLEAN", "date_key": "NUMBER",
                        "const": "VARCHAR", "null_num": "NUMBER", "null_str": "VARCHAR",
                        "null_ts": "TIMESTAMP_NTZ",
                    }
                    col_defs = ", ".join(
                        f'"{col.alias}" {kind_to_type.get(col.kind, "VARCHAR")}'
                        for col in spec.columns
                    )
                    dst_cur.execute(f'CREATE TABLE IF NOT EXISTS "{target}" ({col_defs})')
                    dst_cur.execute(f'TRUNCATE TABLE IF EXISTS "{target}"')
                    logger.info("DM Load: %s -- DDL + truncate done", target)

                # Stream and insert
                logger.debug("DM Load SQL: %s", select_sql)
                src_cur.execute(select_sql)
                col_names = [desc[0] for desc in src_cur.description]
                total_rows = 0
                placeholders = ", ".join(["%s"] * len(col_names))
                ins_sql = f'INSERT INTO "{target}" ({", ".join(col_names)}) VALUES ({placeholders})'

                while True:
                    chunk = src_cur.fetchmany(batch_size)
                    if not chunk:
                        break
                    dst_cur.executemany(ins_sql, chunk)
                    total_rows += len(chunk)
                    logger.info("DM Load: %s -- %d rows so far", target, total_rows)

                src_cur.close()
                dst_cur.close()
            finally:
                # Only close connections we created ourselves
                if own_src:
                    try:
                        src_conn.close()
                    except Exception:
                        pass
                if own_dst:
                    try:
                        dst_conn.close()
                    except Exception:
                        pass

            duration = round(time.time() - t0, 2)
            logger.info("DM Load: %s -- DONE: %d rows in %.1fs", target, total_rows, duration)
            return {
                "status": "loaded",
                "source": source,
                "rows": total_rows,
                "duration_seconds": duration,
            }
        except Exception as exc:
            logger.error("Failed to load %s: %s", target, exc, exc_info=True)
            return {"status": "error", "error": str(exc), "rows": 0}
