from __future__ import annotations

import json
import os
import re
import sys
from pathlib import Path

from ingestion.common.models import Column, Table, table_from_manifest
from ingestion.common.sql import quote_ident, sql_string

from ingestion.common.cloud_lake import client

from .config import LoadConfig, TableTargets
from .runtime_log import upload_runtime_log, write_event


TYPE_MAP = {
    "varchar": "STRING",
    "char": "STRING",
    "text": "STRING",
    "mediumtext": "STRING",
    "longtext": "STRING",
    "datetime": "TIMESTAMP",
    "timestamp": "TIMESTAMP",
    "date": "DATE",
    "json": "JSON",
}


def load_run(config: LoadConfig, manifest_path: str | Path, tables: list[str] | None) -> list[dict]:
    manifest = json.loads(Path(manifest_path).read_text())
    targets = config.resolve_contracts()
    selected = set(tables or [])
    results = []
    skipped = []
    try:
        for item in manifest.get("tables", []):
            table_name = item["table_name"]
            if selected and table_name not in selected:
                continue
            # Only load tables this source has onboarded (a resolved raw target). A manifest
            # may carry tables we don't load — e.g. excluded ones still present in the export.
            if table_name not in targets:
                skipped.append(table_name)
                continue
            results.append(load_table(config, manifest, table_name, targets))
    finally:
        # Upload the run log once, even if a table failed, so a partial run is observable.
        # load_table writes its failure event locally before re-raising, so this final
        # upload captures the crash too.
        upload_runtime_log(config, manifest)
    if skipped:
        print(json.dumps({"skipped_no_target": sorted(skipped)}), file=sys.stderr)
    return results


def load_table(config: LoadConfig, manifest: dict, table_name: str, targets: dict[str, TableTargets]) -> dict:
    table = one_table(manifest, table_name)
    target = targets.get(table.name)
    if target is None or not target.raw_table:
        raise RuntimeError(
            f"no raw target resolved for table {table.name}; "
            "check the load source's raw-snapshot contract"
        )
    # The load stage materializes only the raw layer. Current-state and delta projections
    # belong to the transform layer (dbt/airflow), not here (see ADR 0016).
    raw_asset = target.raw_table
    raw_table = quote_qualified(raw_asset)

    cursor = client(config).cursor()
    try:
        create_raw_table(cursor, raw_table, table)
        # Make the load idempotent and crash-safe before copying. Loading a run_id is
        # defined to converge on exactly one copy of that run's rows however many times it
        # runs, so we always clear this run_id's rows first — that removes the doubling
        # FORCE = TRUE would otherwise cause on a re-run. We also clear orphaned unstamped
        # rows (NULL run_id) left by a crash between COPY and stamp on a prior run, so the
        # stamp below only ever tags this run's freshly-copied rows. Other run_ids are
        # untouched, so multi-snapshot accumulation is preserved. FORCE = TRUE on the COPY
        # is required so the parquet files are re-read after this delete.
        # (Single-writer: not safe for concurrent loads into the same raw table.)
        clear_run_rows(cursor, raw_table, manifest["run_id"])
        copy_snapshot(cursor, raw_table, config, table, manifest)
        stamp_snapshot_metadata(cursor, raw_table, manifest, table)
        loaded_rows = count_rows(cursor, raw_table, manifest["run_id"])
        duplicate_rows = count_duplicate_keys(cursor, raw_table, table, manifest["run_id"])
    except Exception as error:
        # Emit a failure event so a partial run (earlier tables loaded, this one not) is
        # observable in the runtime log, then re-raise to abort the run.
        emit_load_failure_event(config, manifest, table, raw_asset, error)
        raise
    status = "blocked" if duplicate_rows > 0 else "loaded"
    result = {
        "table": table.name,
        "raw_table": raw_asset,
        "loaded_rows": loaded_rows,
        "duplicate_primary_key_rows": duplicate_rows,
        "status": status,
    }
    result["runtime_log"] = write_load_event(
        config=config,
        manifest=manifest,
        table=table,
        raw_table=raw_asset,
        loaded_rows=loaded_rows,
        duplicate_rows=duplicate_rows,
        status=status,
    )
    return result


def one_table(manifest: dict, table_name: str) -> Table:
    summary = manifest_table_summary(manifest, table_name)
    if not summary:
        raise RuntimeError(f"table not found in manifest: {table_name}")
    return table_from_manifest(summary)


def create_raw_table(cursor, raw_table: str, table: Table) -> None:
    source_columns = ",\n  ".join(f"{quote_ident(column.name)} {cloud_lake_type(column)}" for column in table.columns)
    cursor.execute(
        f"""
CREATE TABLE IF NOT EXISTS {raw_table} (
  {source_columns},
  _snapshot_date DATE,
  _snapshot_run_id STRING,
  _snapshot_loaded_at TIMESTAMP,
  _source_database STRING,
  _source_table STRING
)
"""
    )


def cloud_lake_type(column: Column) -> str:
    data_type = column.data_type.lower()
    if data_type in TYPE_MAP:
        return TYPE_MAP[data_type]
    if data_type == "tinyint":
        return "TINYINT"
    if data_type in {"int", "integer"}:
        return "INT"
    if data_type == "bigint":
        return "BIGINT"
    if data_type == "decimal":
        return column.column_type.upper()
    return "STRING"


def clear_run_rows(cursor, raw_table: str, run_id: str) -> None:
    """Delete this run's rows plus any orphaned unstamped rows, making the load idempotent
    and crash-safe. Scoped to this ``run_id`` (plus NULL orphans) so other snapshot runs
    in the same table are preserved."""
    cursor.execute(
        f"""
DELETE FROM {raw_table}
WHERE _snapshot_run_id = {sql_string(run_id)}
   OR _snapshot_run_id IS NULL
"""
    )


def copy_snapshot(cursor, raw_table: str, config: LoadConfig, table: Table, manifest: dict) -> None:
    column_list = ", ".join(quote_ident(column.name) for column in table.columns)
    file_pattern = f".*{re.escape(table.database)}[.]{re.escape(table.name)}[.].*[.]parquet"
    sql = f"""
COPY INTO {raw_table} ({column_list})
FROM {sql_string(manifest["s3_uri"])}
CONNECTION = (
  ACCESS_KEY_ID = {sql_string(config.s3_access_key_id)},
  SECRET_ACCESS_KEY = {sql_string(config.s3_secret_access_key)},
  REGION = {sql_string(config.aws_region)}
)
FILE_FORMAT = (TYPE = PARQUET)
PATTERN = {sql_string(file_pattern)}
COLUMN_MATCH_MODE = CASE_SENSITIVE
FORCE = TRUE
"""
    try:
        cursor.execute(sql)
    except Exception as error:
        raise RuntimeError(sanitize_error(error, config.s3_credential_env_names)) from error


def stamp_snapshot_metadata(cursor, raw_table: str, manifest: dict, table: Table) -> None:
    cursor.execute(
        f"""
UPDATE {raw_table}
SET
  _snapshot_date = {sql_string(manifest["run_date"])}::DATE,
  _snapshot_run_id = {sql_string(manifest["run_id"])},
  _snapshot_loaded_at = now(),
  _source_database = {sql_string(table.database)},
  _source_table = {sql_string(table.name)}
WHERE _snapshot_run_id IS NULL
"""
    )


def count_rows(cursor, raw_table: str, run_id: str) -> int:
    cursor.execute(f"SELECT count(*) FROM {raw_table} WHERE _snapshot_run_id = {sql_string(run_id)}")
    return int(cursor.fetchone()[0])


def count_duplicate_keys(cursor, raw_table: str, table: Table, run_id: str) -> int:
    if not table.primary_key:
        return 0
    key_expr = ", ".join(quote_ident(column) for column in table.primary_key)
    cursor.execute(
        f"""
SELECT coalesce(sum(cnt - 1), 0)
FROM (
  SELECT {key_expr}, count(*) AS cnt
  FROM {raw_table}
  WHERE _snapshot_run_id = {sql_string(run_id)}
  GROUP BY {key_expr}
  HAVING count(*) > 1
) AS dup
"""
    )
    return int(cursor.fetchone()[0])


def write_load_event(
    *,
    config: LoadConfig,
    manifest: dict,
    table: Table,
    raw_table: str,
    loaded_rows: int,
    duplicate_rows: int,
    status: str,
) -> dict[str, str]:
    table_summary = manifest_table_summary(manifest, table.name)
    return write_event(
        config,
        manifest,
        {
            "event_type": "raw_snapshot_loaded",
            "status": status,
            "product_status": manifest.get("product_status", "unknown"),
            "source_database": table.database,
            "source_table": table.name,
            "raw_table": raw_table,
            "primary_key": table.primary_key,
            "watermark_columns": table.watermark_columns,
            "classifications": table.classifications,
            "source_file_count": table_summary.get("file_count", 0),
            "exported_bytes": table_summary.get("total_bytes", 0),
            "loaded_rows": loaded_rows,
            "duplicate_primary_key_rows": duplicate_rows,
            "schema_changed": False,
        },
    )


def emit_load_failure_event(config: LoadConfig, manifest: dict, table: Table, raw_table: str, error: Exception) -> None:
    """Best-effort runtime event for an aborted table load. Never raises: a secondary
    failure here (e.g. the runtime-log upload) must not mask the original load error."""
    try:
        table_summary = manifest_table_summary(manifest, table.name)
        write_event(
            config,
            manifest,
            {
                "event_type": "raw_snapshot_loaded",
                "status": "failed",
                "product_status": manifest.get("product_status", "unknown"),
                "source_database": table.database,
                "source_table": table.name,
                "raw_table": raw_table,
                "error": sanitize_error(error, config.s3_credential_env_names),
                "source_file_count": table_summary.get("file_count", 0),
            },
        )
    except Exception:
        pass


def manifest_table_summary(manifest: dict, table_name: str) -> dict:
    for table in manifest.get("tables", []):
        if table.get("table_name") == table_name:
            return table
    return {}



def quote_qualified(qualified: str) -> str:
    """Quote a ``database.table`` identifier inherited from a contract asset ref."""
    database, _, table = qualified.partition(".")
    return f"{quote_ident(database)}.{quote_ident(table)}"


def sanitize_error(error: Exception, credential_env_names: list[str]) -> str:
    message = str(error)
    for key in credential_env_names:
        value = os.environ.get(key)
        if value:
            message = message.replace(value, f"<{key}>")
    return message