from __future__ import annotations

import json
import os
from datetime import UTC, datetime
from pathlib import Path

from ingestion.common.cloud_lake import client
from ingestion.common.sql import quote_ident, sql_string

from .infer import (
    committed_columns,
    ensure_stage,
    infer_columns,
    proposed_path,
    schema_changed,
    source_contract,
    write_contract,
)
from .storage import StorageAdapter
from .strategy import Increment, IncrementalStrategy

# Provenance every external load stamps, regardless of strategy.
COMMON_PROVENANCE = [("_loaded_at", "TIMESTAMP"), ("_source_uri", "STRING"), ("_source_table", "STRING")]


class SchemaDriftError(RuntimeError):
    """Raised when the live inferred schema no longer matches the committed source contract.
    The load aborts and a proposed contract is written for review (a contract update
    process) rather than the committed contract being silently overwritten."""


def external_load_run(config, storage: StorageAdapter, strategy: IncrementalStrategy, reload_key: int | None = None) -> list[dict]:
    """Incrementally load an externally-produced bucket into the raw layer.

    Lists the bucket, splits it into increments via ``strategy``, gates on schema drift
    (routing changes to a proposed contract), and loads only increments whose cursor is
    newer than the max already in the raw table — stamping the cursor onto every row."""
    cursor = client(config).cursor()
    ensure_stage(cursor, storage, config)

    objects = storage.list_objects(config.external_uri)
    increments = strategy.increments(objects)
    if not increments:
        raise RuntimeError(f"no loadable objects under {config.external_uri} for strategy {strategy.name!r}")

    live = infer_columns(cursor, config.stage_name, sample_object(increments))
    committed = committed_columns(config)
    if committed is None:
        raise RuntimeError(f"no committed source contract; run `infer-source` and commit it before loading ({config.external_table})")
    if schema_changed(live, committed):
        path = write_contract(proposed_path(config), source_contract(config, live))
        write_event(config, {"event_type": "schema_drift", "status": "blocked",
                             "source_table": config.external_table, "strategy": strategy.name,
                             "proposed_contract": str(path)})
        raise SchemaDriftError(
            f"inferred schema for {config.external_table} differs from the committed source contract; "
            f"wrote a proposed contract to {path}. Review and promote it (contract update process) before loading."
        )

    target = config.resolve_contracts().get(config.external_table)
    if target is None or not target.raw_table:
        raise RuntimeError(f"no raw target resolved for {config.external_table}; run `generate-contracts` first")
    raw_table = quote_qualified(target.raw_table)
    data_columns = [c["name"] for c in live]
    create_raw_table(cursor, raw_table, live, strategy)

    max_wm = current_max(cursor, raw_table, strategy.watermark_column)
    selected = select_increments(increments, max_wm, config.delta, reload_key)
    relative = relativize(config.external_uri, [name for inc in increments for name in inc.files])
    results = []
    for increment in selected:
        results.append(load_increment(cursor, config, storage, strategy, raw_table, target.raw_table, increment, data_columns, relative))
    if not selected:
        write_event(config, {"event_type": "raw_snapshot_loaded", "status": "skipped_not_newer",
                             "source_table": config.external_table, "strategy": strategy.name,
                             "max_loaded_watermark": max_wm})
    return results


def select_increments(increments: list[Increment], max_wm: int | None, delta: bool, reload_key: int | None) -> list[Increment]:
    """Increments to load, ascending by cursor. Delta on: only those strictly newer than the
    max already loaded (all of them on first load, max None). ``reload_key`` force-includes
    one; delta off loads all."""
    return [
        inc for inc in sorted(increments, key=lambda i: i.key)
        if inc.key == reload_key or not delta or max_wm is None or inc.key > max_wm
    ]


def load_increment(cursor, config, storage, strategy, raw_table, raw_asset, increment, data_columns, relative) -> dict:
    selector = strategy.clear_predicate(increment)
    try:
        clear_increment(cursor, raw_table, strategy.watermark_column, selector)
        copy_increment(cursor, config, storage, raw_table, data_columns, increment, relative)
        stamp_increment(cursor, config, strategy, raw_table, increment)
        loaded_rows = count_rows(cursor, raw_table, selector)
        duplicate_rows = count_duplicate_keys(cursor, raw_table, config.external_primary_key, selector)
    except Exception as error:
        write_event(config, {"event_type": "raw_snapshot_loaded", "status": "failed", "strategy": strategy.name,
                             "increment": increment.label, "source_table": config.external_table, "raw_table": raw_asset,
                             "error": sanitize_error(error, storage.credential_env_names())})
        raise
    status = "blocked" if duplicate_rows > 0 else "loaded"
    result = {
        "table": config.external_table,
        "raw_table": raw_asset,
        "strategy": strategy.name,
        "increment": increment.label,
        "cursor": increment.key,
        "loaded_rows": loaded_rows,
        "duplicate_primary_key_rows": duplicate_rows,
        "status": status,
    }
    write_event(config, {"event_type": "raw_snapshot_loaded", **result})
    return result


def sample_object(increments: list[Increment]) -> str:
    """Pick one object's basename to infer the schema from (the newest increment's first file)."""
    newest = max(increments, key=lambda i: i.key)
    return Path(sorted(newest.files)[0]).name


def relativize(uri: str, names: list[str]) -> dict[str, str]:
    """Map each full object name to its path relative to the stage URL (what COPY FILES wants)."""
    prefix = uri.split("://", 1)[1].split("/", 1)[1] if "/" in uri.split("://", 1)[1] else ""
    return {name: name[len(prefix):] if name.startswith(prefix) else name for name in names}


def create_raw_table(cursor, raw_table: str, columns: list[dict], strategy: IncrementalStrategy) -> None:
    cols = [(c["name"], c["type"]) for c in columns] + strategy.provenance_columns() + COMMON_PROVENANCE
    body = ",\n  ".join(f"{quote_ident(name)} {type_}" for name, type_ in cols)
    cursor.execute(f"CREATE TABLE IF NOT EXISTS {raw_table} (\n  {body}\n)")


def clear_increment(cursor, raw_table: str, watermark_column: str, selector: str) -> None:
    """Drop this increment's rows (idempotent reload) plus crash-orphaned unstamped rows;
    other increments are untouched, preserving accumulation."""
    cursor.execute(f"DELETE FROM {raw_table} WHERE ({selector}) OR {quote_ident(watermark_column)} IS NULL")


def copy_increment(cursor, config, storage, raw_table, data_columns, increment, relative) -> None:
    column_list = ", ".join(quote_ident(name) for name in data_columns)
    files = ", ".join(sql_string(relative[name]) for name in increment.files)
    sql = f"""
COPY INTO {raw_table} ({column_list})
FROM '@{config.stage_name}'
FILES = ({files})
FILE_FORMAT = (TYPE = PARQUET)
COLUMN_MATCH_MODE = CASE_SENSITIVE
FORCE = TRUE
"""
    try:
        cursor.execute(sql)
    except Exception as error:
        raise RuntimeError(sanitize_error(error, storage.credential_env_names())) from error


def stamp_increment(cursor, config, strategy, raw_table, increment) -> None:
    assignments = {
        **strategy.stamp_assignments(increment),
        "_loaded_at": "now()",
        "_source_uri": sql_string(config.external_uri),
        "_source_table": sql_string(config.external_table),
    }
    set_clause = ",\n  ".join(f"{quote_ident(col)} = {value}" for col, value in assignments.items())
    cursor.execute(
        f"UPDATE {raw_table}\nSET\n  {set_clause}\nWHERE {quote_ident(strategy.watermark_column)} IS NULL"
    )


def current_max(cursor, raw_table: str, watermark_column: str) -> int | None:
    cursor.execute(f"SELECT max({quote_ident(watermark_column)}) FROM {raw_table}")
    value = cursor.fetchone()[0]
    return int(value) if value is not None else None


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


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


def quote_qualified(qualified: str) -> str:
    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


def write_event(config, event: dict) -> None:
    """Append one run event to a local JSONL log (external runs have no manifest and we only
    have read access to the source bucket, so the log stays local)."""
    run_dir = config.root / "runs" / datetime.now(UTC).strftime("%Y-%m-%d")
    run_dir.mkdir(parents=True, exist_ok=True)
    record = {"source_name": config.name, "created_at": datetime.now(UTC).isoformat(), **event}
    with (run_dir / "events.jsonl").open("a") as file:
        file.write(json.dumps(record, sort_keys=True) + "\n")