from __future__ import annotations

import json
import os
import re
import subprocess
import tempfile
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

import yaml

from orchestration.airflow.engine.contract_assets import (
    AirflowAssetSpec,
    CompiledAssetGraph,
    CompiledContractAsset,
    compile_repo_contract_assets,
)


REPO_ROOT = Path(__file__).resolve().parents[3]


@dataclass(frozen=True)
class LoadAssetCommand:
    executable_path: Path
    loader_module: str
    source_path: Path
    mode: str
    tables: tuple[str, ...]
    reload_key: int | None
    env: dict[str, str]
    cwd: Path

    def argv(self, *, manifest_path: str | Path | None = None) -> list[str]:
        argv = [
            str(self.executable_path),
            self.loader_module,
            "--source",
            str(self.source_path),
            "load-run",
        ]
        if self.mode == "manifest":
            if manifest_path is None:
                raise ValueError("manifest load commands require a manifest path")
            argv.extend(["--manifest", str(manifest_path)])
            if self.tables:
                argv.extend(["--tables", ",".join(self.tables)])
            return argv
        if self.mode == "external_incremental":
            if self.reload_key is not None:
                argv.extend(["--reload-key", str(self.reload_key)])
            return argv
        raise ValueError(f"unsupported load asset mode: {self.mode!r}")


@dataclass(frozen=True)
class LoadAssetUnitConfig:
    name: str
    mode: str
    dag_id: str
    task_id: str
    owner: str
    start_date: str
    tags: tuple[str, ...]
    catchup: bool
    max_active_runs: int
    schedule: str | tuple[AirflowAssetSpec, ...]
    command: LoadAssetCommand
    inlets: tuple[AirflowAssetSpec, ...]
    outlets: tuple[AirflowAssetSpec, ...]
    manifest_asset: AirflowAssetSpec | None
    manifest_path: Path | None
    manifest_path_env: str | None
    manifest_event_extra_keys: tuple[str, ...]
    runtime_repo_root: Path
    runtime_workspace_root: Path


@dataclass(frozen=True)
class LoadAssetDagConfig:
    units: tuple[LoadAssetUnitConfig, ...]

    def by_name(self) -> dict[str, LoadAssetUnitConfig]:
        return {unit.name: unit for unit in self.units}


def load_load_asset_config(
    config_path: Path,
    *,
    repo_root: Path = REPO_ROOT,
    config_key: str = "load_asset_units",
) -> LoadAssetDagConfig:
    source_config = yaml.safe_load(config_path.read_text()) or {}
    unit_configs = tuple(source_config[config_key])
    graph = compile_repo_contract_assets(repo_root)
    runtime_root = _runtime_repo_root(repo_root)
    runtime_workspace_root = _runtime_load_workspace(runtime_root)

    return LoadAssetDagConfig(
        units=tuple(
            _load_asset_unit_config(
                unit_config,
                graph=graph,
                repo_root=repo_root,
                runtime_root=runtime_root,
                runtime_workspace_root=runtime_workspace_root,
            )
            for unit_config in unit_configs
        )
    )


def build_airflow_asset(spec: AirflowAssetSpec) -> Any:
    from airflow.sdk import Asset

    return Asset(
        uri=spec.uri,
        name=spec.name,
        group=spec.group,
        extra=dict(spec.extra),
    )


def resolve_manifest_path(
    unit: LoadAssetUnitConfig,
    *,
    context: dict[str, Any] | None = None,
    env: dict[str, str] | None = None,
) -> Path:
    if unit.mode != "manifest":
        raise ValueError(f"load asset unit {unit.name!r} is not manifest-driven")

    reference: str | Path | None = unit.manifest_path
    if reference is None and context is not None:
        reference = _manifest_reference_from_context(unit, context)

    env_values = os.environ if env is None else env
    if reference is None and unit.manifest_path_env:
        reference = env_values.get(unit.manifest_path_env)

    if reference is None:
        raise RuntimeError(
            f"manifest-driven load asset unit {unit.name!r} needs a manifest path. "
            "Pass it in the triggering asset event metadata or set "
            f"{unit.manifest_path_env or 'the configured manifest path env var'}."
        )

    return _materialize_manifest_reference(
        reference,
        source_repo_root=unit.runtime_repo_root,
        workspace_repo_root=unit.runtime_workspace_root,
    )


def execute_load_asset_command(
    command: LoadAssetCommand,
    *,
    output_assets: tuple[AirflowAssetSpec, ...],
    manifest_path: str | Path | None = None,
) -> dict[str, object]:
    argv = command.argv(manifest_path=manifest_path)
    command_env = {**os.environ, **command.env}
    try:
        completed = _run_command(
            argv,
            cwd=command.cwd,
            env=command_env,
        )
    except RuntimeError as error:
        action_request_path = _write_contract_update_pr_request(
            command,
            error_message=str(error),
            env=command_env,
        )
        if action_request_path:
            raise RuntimeError(
                f"{error}\ncontract-update PR action request: {action_request_path}"
            ) from error
        raise
    payload = json.loads(completed.stdout or "{}")
    loads = _validated_load_results(
        payload,
        mode=command.mode,
        output_assets=output_assets,
    )
    return {
        "outputs": [
            {
                "asset_name": asset.name,
                "asset_uri": asset.uri,
                "asset_group": asset.group,
            }
            for asset in output_assets
        ],
        "loader": {**payload, "loads": loads},
    }


def _validated_load_results(
    payload: dict[str, Any],
    *,
    mode: str,
    output_assets: tuple[AirflowAssetSpec, ...],
) -> list[dict[str, Any]]:
    loads = payload.get("loads")
    if not isinstance(loads, list):
        raise RuntimeError("load command did not return a JSON object with a loads array")
    if not all(isinstance(item, dict) for item in loads):
        raise RuntimeError("load command returned non-object items in its loads array")

    if mode == "external_incremental" and not loads:
        _raise_airflow_skip("external-incremental poll loaded no new increments")

    if mode == "manifest":
        loaded_raw_tables = {str(item.get("raw_table")) for item in loads if item.get("raw_table")}
        expected_raw_tables = {
            asset.extra["asset_ref"]
            for asset in output_assets
            if asset.extra.get("asset_ref")
        }
        if len(expected_raw_tables) != len(output_assets):
            raise RuntimeError("manifest load output assets must carry contract asset_ref metadata")
        missing = sorted(expected_raw_tables - loaded_raw_tables)
        if missing:
            raise RuntimeError(
                "manifest load did not materialize every declared outlet raw table: "
                + ", ".join(missing)
            )

    return loads


def _raise_airflow_skip(message: str) -> None:
    from airflow.sdk.exceptions import AirflowSkipException

    raise AirflowSkipException(message)


def _write_contract_update_pr_request(
    command: LoadAssetCommand,
    *,
    error_message: str,
    env: dict[str, str],
) -> Path | None:
    if command.mode != "external_incremental":
        return None
    proposed_contract_path = _proposed_contract_path_from_error(error_message)
    if proposed_contract_path is None:
        return None

    source_name = _source_name(command.source_path)
    request_id = _action_request_id(source_name)
    request = {
        "request_id": request_id,
        "status": "accepted",
        "action": "open_pr",
        "risk_tier": "controlled_operational",
        "reason": (
            "External-incremental schema drift blocked the load; review and "
            "promote the proposed source contract through Git."
        ),
        "parameters": {
            "source_name": source_name,
            "source_path": str(command.source_path),
            "loader_module": command.loader_module,
            "proposed_contract_path": proposed_contract_path,
            "promote_to_path": _promoted_contract_path(proposed_contract_path),
            "branch_name": f"codex/contract-update-{source_name}-schema-drift",
            "commit_message": f"Promote {source_name} external source contract drift",
            "pr_title": f"[contract-update] Review {source_name} source schema drift",
            "pr_body": (
                "Schema drift was detected by the external-incremental load. "
                "Review the proposed source contract, promote it to "
                "`contracts/upstream/` if correct, regenerate the raw contract, "
                "and rerun the load."
            ),
        },
    }

    action_dir = _source_run_artifact_root(command, env) / "action-requests"
    action_dir.mkdir(parents=True, exist_ok=True)
    path = action_dir / f"{request_id}.json"
    path.write_text(json.dumps(request, indent=2, sort_keys=True) + "\n")
    return path


def _proposed_contract_path_from_error(error_message: str) -> str | None:
    match = re.search(r"(?P<path>\S*contracts/proposed/\S+)", error_message)
    if match is None:
        return None
    return match.group("path").rstrip(".,;:)")


def _promoted_contract_path(proposed_contract_path: str) -> str:
    return proposed_contract_path.replace(
        "contracts/proposed/",
        "contracts/upstream/",
        1,
    )


def _source_name(source_path: Path) -> str:
    return source_path.parent.name or "unknown-source"


def _action_request_id(source_name: str) -> str:
    timestamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
    safe_source = re.sub(r"[^A-Za-z0-9_.-]+", "-", source_name).strip("-") or "source"
    return f"actreq-{timestamp}-{safe_source}-contract-update-pr"


def _source_run_artifact_root(command: LoadAssetCommand, env: dict[str, str]) -> Path:
    source_path = command.source_path
    if source_path.is_absolute():
        return source_path.parent / "runs"
    workspace = Path(
        env.get("DATALAKE_LOAD_WORKSPACE")
        or env.get("DATALAKE_AIRFLOW_REPO_ROOT")
        or command.cwd
    )
    return workspace / source_path.parent / "runs"


def _failed_command_message(argv: list[str], error: subprocess.CalledProcessError) -> str:
    stdout = _redacted_excerpt(error.stdout)
    stderr = _redacted_excerpt(error.stderr)
    lines = [
        f"load command failed with exit code {error.returncode}: {' '.join(argv)}",
    ]
    if stdout:
        lines.append(f"stdout:\n{stdout}")
    if stderr:
        lines.append(f"stderr:\n{stderr}")
    return "\n".join(lines)


def _redacted_excerpt(value: str | bytes | None, limit: int = 4000) -> str:
    if value is None:
        return ""
    text = value.decode() if isinstance(value, bytes) else value
    for key, secret in os.environ.items():
        if not secret or len(secret) < 8:
            continue
        if any(marker in key.upper() for marker in ("SECRET", "TOKEN", "PASSWORD", "CREDENTIAL", "DSN", "KEY")):
            text = text.replace(secret, f"<{key}>")
    text = text.strip()
    if len(text) > limit:
        return text[:limit] + "\n... truncated ..."
    return text


def _load_asset_unit_config(
    unit_config: dict[str, Any],
    *,
    graph: CompiledAssetGraph,
    repo_root: Path,
    runtime_root: Path,
    runtime_workspace_root: Path,
) -> LoadAssetUnitConfig:
    mode = unit_config["mode"]
    source_path = _resolve_runtime_path(repo_root, unit_config["source_path"])
    runtime_source_path = Path(unit_config["source_path"])
    source_data = yaml.safe_load(source_path.read_text()) or {}
    _validate_mode_matches_source(mode, source_data, source_path)

    produced_by_ref = unit_config["produced_by_ref"]
    produced_assets = _produced_assets_for_source(
        graph,
        source_path=source_path,
        repo_root=repo_root,
        produced_by_ref=produced_by_ref,
        output_contracts=tuple(unit_config.get("output_contracts", ())),
    )
    outlets = tuple(asset.airflow for asset in produced_assets)

    manifest_config = unit_config.get("manifest", {})
    manifest_asset = (
        _airflow_asset_spec(manifest_config["asset"])
        if mode == "manifest"
        else None
    )
    manifest_path = (
        _resolve_runtime_path(runtime_workspace_root, manifest_config["path"])
        if manifest_config.get("path")
        else None
    )
    manifest_path_env = manifest_config.get("path_env")
    manifest_event_extra_keys = tuple(
        manifest_config.get("event_extra_keys", ("manifest_path", "manifest_uri"))
    )

    contract_inlets = _shared_inlets_for_outputs(graph, produced_assets)
    inlets = _dedupe_assets(
        (
            *((manifest_asset,) if manifest_asset else ()),
            *contract_inlets,
        )
    )

    if mode == "manifest":
        if manifest_asset is None:
            raise ValueError(f"manifest load asset unit {unit_config['name']!r} must declare manifest.asset")
        schedule: str | tuple[AirflowAssetSpec, ...] = (manifest_asset,)
    elif mode == "external_incremental":
        schedule = unit_config["poll_schedule"]
    else:
        raise ValueError(f"unsupported load asset mode: {mode!r}")

    runner = unit_config.get("runner", {})
    command_env = {
        "DATALAKE_AIRFLOW_REPO_ROOT": str(runtime_root),
        **{str(key): str(value) for key, value in runner.get("env", {}).items()},
    }
    if runner.get("uv_project_environment"):
        command_env["UV_PROJECT_ENVIRONMENT"] = str(runner["uv_project_environment"])

    return LoadAssetUnitConfig(
        name=unit_config["name"],
        mode=mode,
        dag_id=unit_config["dag_id"],
        task_id=unit_config["task_id"],
        owner=unit_config["owner"],
        start_date=unit_config["start_date"],
        tags=tuple(unit_config["tags"]),
        catchup=bool(unit_config.get("catchup", False)),
        max_active_runs=int(unit_config.get("max_active_runs", 1)),
        schedule=schedule,
        command=LoadAssetCommand(
            executable_path=_resolve_runtime_path(runtime_root, runner["executable_path"]),
            loader_module=unit_config["loader_module"],
            source_path=runtime_source_path,
            mode=mode,
            tables=tuple(unit_config.get("tables", ())),
            reload_key=unit_config.get("reload_key"),
            env=command_env,
            cwd=runtime_root,
        ),
        inlets=inlets,
        outlets=outlets,
        manifest_asset=manifest_asset,
        manifest_path=manifest_path,
        manifest_path_env=manifest_path_env,
        manifest_event_extra_keys=manifest_event_extra_keys,
        runtime_repo_root=runtime_root,
        runtime_workspace_root=runtime_workspace_root,
    )


def _produced_assets_for_source(
    graph: CompiledAssetGraph,
    *,
    source_path: Path,
    repo_root: Path,
    produced_by_ref: str,
    output_contracts: tuple[str, ...],
) -> tuple[CompiledContractAsset, ...]:
    generated_prefix = _relative_path(source_path.parent / "contracts" / "generated", repo_root) + "/"
    candidates = {
        asset.contract_id: asset
        for asset in graph.assets.values()
        if asset.contract_path.startswith(generated_prefix)
    }
    if output_contracts:
        missing = sorted(contract_id for contract_id in output_contracts if contract_id not in candidates)
        if missing:
            raise ValueError(
                f"load source {source_path} is missing output contract(s): {', '.join(missing)}"
            )
        candidates = {contract_id: candidates[contract_id] for contract_id in output_contracts}
    if not candidates:
        raise ValueError(f"load source {source_path} has no generated raw-snapshot contracts")

    for asset in candidates.values():
        if asset.contract_kind != "product" or asset.layer != "raw":
            raise ValueError(
                f"load output {asset.contract_id!r} must be a raw product contract; "
                f"found kind={asset.contract_kind!r}, layer={asset.layer!r}"
            )
        if asset.produced_by_kind != "python" or asset.produced_by_ref != produced_by_ref:
            raise ValueError(
                f"load output {asset.contract_id!r} producer does not match "
                f"{produced_by_ref!r}: {asset.produced_by_kind!r}/{asset.produced_by_ref!r}"
            )

    return tuple(candidates[contract_id] for contract_id in sorted(candidates))


def _shared_inlets_for_outputs(
    graph: CompiledAssetGraph,
    produced_assets: tuple[CompiledContractAsset, ...],
) -> tuple[AirflowAssetSpec, ...]:
    inlet_by_name: dict[str, AirflowAssetSpec] = {}
    for asset in produced_assets:
        for edge in graph.edges:
            if edge.downstream_ref != asset.contract_id:
                continue
            upstream_asset = graph.assets.get(edge.upstream_ref)
            if upstream_asset:
                inlet_by_name[upstream_asset.airflow.name] = upstream_asset.airflow
                continue
            external_asset = graph.external_assets.get(edge.upstream_ref)
            if external_asset:
                inlet_by_name[external_asset.airflow.name] = external_asset.airflow
    return tuple(inlet_by_name[name] for name in sorted(inlet_by_name))


def _airflow_asset_spec(config: dict[str, Any]) -> AirflowAssetSpec:
    return AirflowAssetSpec(
        name=config["name"],
        uri=config["uri"],
        group=config["group"],
        extra={str(key): str(value) for key, value in config.get("extra", {}).items()},
    )


def _dedupe_assets(assets: tuple[AirflowAssetSpec, ...]) -> tuple[AirflowAssetSpec, ...]:
    by_name = {asset.name: asset for asset in assets}
    return tuple(by_name[name] for name in sorted(by_name))


def _manifest_reference_from_context(unit: LoadAssetUnitConfig, context: dict[str, Any]) -> str | None:
    for event in _manifest_events_from_context(unit, context):
        extra = _event_extra(event)
        for key in unit.manifest_event_extra_keys:
            value = extra.get(key)
            if value:
                return str(value)
    return None


def _manifest_events_from_context(unit: LoadAssetUnitConfig, context: dict[str, Any]) -> list[Any]:
    manifest_asset = unit.manifest_asset
    if manifest_asset is None:
        return []

    candidate_events: list[Any] = []
    fallback_events: list[Any] = []
    for key in ("triggering_asset_events", "triggering_dataset_events"):
        mapping = context.get(key)
        if not isinstance(mapping, dict):
            continue
        for asset_key, events in mapping.items():
            event_list = list(events if isinstance(events, list | tuple) else (events,))
            fallback_events.extend(event_list)
            key_text = str(asset_key)
            if asset_key in {manifest_asset.name, manifest_asset.uri}:
                candidate_events.extend(event_list)
            elif manifest_asset.name in key_text or manifest_asset.uri in key_text:
                candidate_events.extend(event_list)
    return candidate_events or fallback_events


def _event_extra(event: Any) -> dict[str, Any]:
    if isinstance(event, dict):
        extra = event.get("extra", {})
    else:
        extra = getattr(event, "extra", {})
    return extra if isinstance(extra, dict) else {}


def _materialize_manifest_reference(
    reference: str | Path,
    *,
    source_repo_root: Path,
    workspace_repo_root: Path,
) -> Path:
    text = str(reference)
    path = Path(text)
    if text.startswith("s3://") or text.startswith("gs://"):
        destination_dir = Path(tempfile.mkdtemp(prefix="datalake-load-manifest-"))
        destination = destination_dir / "manifest.json"
        _download_object(text, destination)
        return destination
    if path.is_absolute():
        try:
            return workspace_repo_root / path.relative_to(source_repo_root)
        except ValueError:
            pass
        return path
    return workspace_repo_root / path


def _download_object(uri: str, destination: Path) -> None:
    destination.parent.mkdir(parents=True, exist_ok=True)
    if uri.startswith("s3://"):
        command = ["aws", "s3", "cp", uri, str(destination), "--only-show-errors"]
    elif uri.startswith("gs://"):
        command = ["gcloud", "storage", "cp", uri, str(destination)]
    else:
        raise ValueError(f"unsupported manifest URI scheme: {uri!r}")
    _run_command(command)


def _run_command(
    argv: list[str],
    *,
    cwd: Path | None = None,
    env: dict[str, str] | None = None,
) -> subprocess.CompletedProcess[str]:
    try:
        return subprocess.run(
            argv,
            check=True,
            capture_output=True,
            text=True,
            cwd=cwd,
            env=env,
        )
    except subprocess.CalledProcessError as error:
        raise RuntimeError(_failed_command_message(argv, error)) from error
    except OSError as error:
        raise RuntimeError(f"command failed to start: {' '.join(argv)}\n{error}") from error


def _validate_mode_matches_source(mode: str, source_data: dict[str, Any], source_path: Path) -> None:
    is_external = "external" in source_data
    if mode == "manifest" and is_external:
        raise ValueError(f"manifest load asset unit points at an external source: {source_path}")
    if mode == "external_incremental" and not is_external:
        raise ValueError(f"external load asset unit points at a manifest-driven source: {source_path}")


def _runtime_repo_root(repo_root: Path) -> Path:
    override = os.environ.get("DATALAKE_AIRFLOW_REPO_ROOT")
    if override:
        return Path(override)
    return repo_root


def _runtime_load_workspace(runtime_root: Path) -> Path:
    override = os.environ.get("DATALAKE_LOAD_WORKSPACE")
    if override:
        return Path(override)
    return runtime_root


def _resolve_runtime_path(runtime_root: Path, value: str) -> Path:
    path = Path(value)
    if path.is_absolute():
        return path
    return runtime_root / path


def _relative_path(path: Path, repo_root: Path) -> str:
    try:
        return str(path.relative_to(repo_root))
    except ValueError:
        return str(path)