from __future__ import annotations

import json
import os
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import yaml

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


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


@dataclass(frozen=True)
class ExtractAssetCommand:
    executable_path: Path
    extractor_module: str
    source_path: Path
    mode: str
    tables: tuple[str, ...]
    env: dict[str, str]
    cwd: Path

    def argv(
        self,
        command: str,
        *,
        run_date: str | None = None,
        run_id: str | None = None,
        export_id: str | None = None,
    ) -> list[str]:
        argv = [
            str(self.executable_path),
            self.extractor_module,
            "--source",
            str(self.source_path),
            command,
        ]
        if command in {"export", "start-export"}:
            if run_date:
                argv.extend(["--run-date", run_date])
            if run_id:
                argv.extend(["--run-id", run_id])
            if command == "export" and self.tables:
                argv.extend(["--tables", ",".join(self.tables)])
            return argv
        if command == "write-manifest":
            if not run_date or not run_id:
                raise ValueError("write-manifest requires run_date and run_id")
            argv.extend(["--run-date", run_date, "--run-id", run_id])
            if self.mode == "tidbcloud_deferrable":
                if not export_id:
                    raise ValueError("TiDB Cloud write-manifest requires export_id")
                argv.extend(["--export-id", export_id])
            return argv
        if command == "export-status":
            if not export_id:
                raise ValueError("export-status requires export_id")
            argv.extend(["--export-id", export_id])
            return argv
        if command == "generate-contracts":
            return argv
        raise ValueError(f"unsupported extract command: {command!r}")

    def to_trigger_kwargs(self) -> dict[str, object]:
        return {
            "executable_path": str(self.executable_path),
            "extractor_module": self.extractor_module,
            "source_path": str(self.source_path),
            "mode": self.mode,
            "tables": list(self.tables),
            "env": dict(self.env),
            "cwd": str(self.cwd),
        }

    @classmethod
    def from_trigger_kwargs(cls, payload: dict[str, object]) -> ExtractAssetCommand:
        return cls(
            executable_path=Path(str(payload["executable_path"])),
            extractor_module=str(payload["extractor_module"]),
            source_path=Path(str(payload["source_path"])),
            mode=str(payload["mode"]),
            tables=tuple(str(item) for item in payload.get("tables", [])),
            env={str(key): str(value) for key, value in dict(payload.get("env", {})).items()},
            cwd=Path(str(payload["cwd"])),
        )


@dataclass(frozen=True)
class ExtractAssetUnitConfig:
    name: str
    mode: str
    dag_id: str
    owner: str
    start_date: str
    tags: tuple[str, ...]
    catchup: bool
    max_active_runs: int
    schedule: str | None
    command: ExtractAssetCommand
    manifest_asset: AirflowAssetSpec
    source_outlets: tuple[AirflowAssetSpec, ...]
    generate_contracts_task_id: str
    export_task_id: str
    write_manifest_task_id: str
    poll_interval_seconds: int
    timeout_seconds: int


@dataclass(frozen=True)
class ExtractAssetDagConfig:
    units: tuple[ExtractAssetUnitConfig, ...]

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


def load_extract_asset_config(
    config_path: Path,
    *,
    repo_root: Path = REPO_ROOT,
    config_key: str = "extract_asset_units",
) -> ExtractAssetDagConfig:
    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)

    return ExtractAssetDagConfig(
        units=tuple(
            _extract_asset_unit_config(
                unit_config,
                graph=graph,
                repo_root=repo_root,
                runtime_root=runtime_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 execute_extract_command(
    command: ExtractAssetCommand,
    subcommand: str,
    *,
    run_date: str | None = None,
    run_id: str | None = None,
    export_id: str | None = None,
) -> dict[str, Any]:
    completed = _run_command(
        command.argv(subcommand, run_date=run_date, run_id=run_id, export_id=export_id),
        cwd=command.cwd,
        env={**os.environ, **command.env},
    )
    payload = json.loads(completed.stdout or "{}")
    return _normalize_extract_payload(payload, command=command)


def manifest_event_extra(
    *,
    manifest_payload: dict[str, Any],
    fallback_payload: dict[str, Any],
) -> dict[str, str]:
    manifest_path = str(manifest_payload["manifest"])
    extra = {
        "manifest_path": manifest_path,
        "run_date": str(fallback_payload.get("run_date", "")),
        "run_id": str(fallback_payload.get("run_id", "")),
    }
    path = Path(manifest_path)
    if path.exists():
        manifest = json.loads(path.read_text())
        for key in ("manifest_uri", "run_date", "run_id", "source_name"):
            if manifest.get(key):
                extra[key] = str(manifest[key])
    return extra


def _extract_asset_unit_config(
    unit_config: dict[str, Any],
    *,
    graph: CompiledAssetGraph,
    repo_root: Path,
    runtime_root: Path,
) -> ExtractAssetUnitConfig:
    mode = unit_config["mode"]
    source_path = _resolve_runtime_path(repo_root, unit_config["source_path"])
    source_data = yaml.safe_load(source_path.read_text()) or {}
    source_outlets = _source_outlets_for_source(
        graph,
        source_path=source_path,
        repo_root=repo_root,
        produced_by_ref=unit_config["produced_by_ref"],
        output_contracts=tuple(unit_config.get("output_contracts", ())),
    )
    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"])
    if runner.get("workspace_env"):
        command_env["DATALAKE_EXTRACT_WORKSPACE"] = str(runner["workspace_env"])

    return ExtractAssetUnitConfig(
        name=unit_config["name"],
        mode=mode,
        dag_id=unit_config["dag_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_for_unit(unit_config, source_data),
        command=ExtractAssetCommand(
            executable_path=_resolve_runtime_path(runtime_root, runner["executable_path"]),
            extractor_module=unit_config["extractor_module"],
            source_path=Path(unit_config["source_path"]),
            mode=mode,
            tables=tuple(unit_config.get("tables", ())),
            env=command_env,
            cwd=runtime_root,
        ),
        manifest_asset=_airflow_asset_spec(unit_config["manifest"]["asset"]),
        source_outlets=source_outlets,
        generate_contracts_task_id=unit_config.get("generate_contracts_task_id", "generate_contracts"),
        export_task_id=unit_config.get("export_task_id", "export"),
        write_manifest_task_id=unit_config.get("write_manifest_task_id", "write_manifest"),
        poll_interval_seconds=int(unit_config.get("poll_interval_seconds", 300)),
        timeout_seconds=int(unit_config.get("timeout_seconds", 7200)),
    )


def _source_outlets_for_source(
    graph: CompiledAssetGraph,
    *,
    source_path: Path,
    repo_root: Path,
    produced_by_ref: str,
    output_contracts: tuple[str, ...],
) -> tuple[AirflowAssetSpec, ...]:
    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"extract 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"extract source {source_path} has no generated source contracts")

    for asset in candidates.values():
        if asset.contract_kind != "source":
            raise ValueError(f"extract output {asset.contract_id!r} must be a source contract")
        if asset.produced_by_kind != "python" or asset.produced_by_ref != produced_by_ref:
            raise ValueError(
                f"extract 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].airflow for contract_id in sorted(candidates))


def _normalize_extract_payload(
    payload: dict[str, Any],
    *,
    command: ExtractAssetCommand,
) -> dict[str, Any]:
    manifest = payload.get("manifest")
    if not isinstance(manifest, str) or not manifest:
        return payload

    path = Path(manifest)
    if not path.is_absolute():
        path = _runtime_extract_workspace(command) / path
    return {**payload, "manifest": str(path)}


def _runtime_extract_workspace(command: ExtractAssetCommand) -> Path:
    for env_key in ("DATALAKE_EXTRACT_WORKSPACE", "DATALAKE_AIRFLOW_REPO_ROOT"):
        value = command.env.get(env_key)
        if value:
            return Path(value)
    return command.cwd


def _schedule_for_unit(unit_config: dict[str, Any], source_data: dict[str, Any]) -> str | None:
    if "schedule" in unit_config:
        return unit_config["schedule"]
    return source_data.get("schedule", {}).get("cron")


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 _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 _failed_command_message(argv: list[str], error: subprocess.CalledProcessError) -> str:
    stdout = _redacted_excerpt(error.stdout)
    stderr = _redacted_excerpt(error.stderr)
    lines = [f"extract 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 _runtime_repo_root(repo_root: Path) -> Path:
    override = os.environ.get("DATALAKE_AIRFLOW_REPO_ROOT")
    if override:
        return Path(override)
    return repo_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)