from __future__ import annotations

import importlib
import json
import os
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable

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 PythonAssetJobConfig:
    dag_id: str
    task_id: str
    owner: str
    start_date: str
    tags: tuple[str, ...]
    catchup: bool
    max_active_runs: int
    callable_path: str
    runner_executable_path: Path
    produced_by_ref: str
    run_kwargs: dict[str, Any]
    inlets: tuple[AirflowAssetSpec, ...]
    outlets: tuple[AirflowAssetSpec, ...]
    result_contract_order: tuple[str, ...]


class PythonAssetJobSkipped(RuntimeError):
    pass


def load_python_asset_job_config(
    config_path: Path,
    *,
    repo_root: Path = REPO_ROOT,
    config_key: str = "python_asset_job",
) -> PythonAssetJobConfig:
    source_config = yaml.safe_load(config_path.read_text()) or {}
    job_config = source_config[config_key]
    runtime_root = _runtime_repo_root(repo_root)

    graph = compile_repo_contract_assets(repo_root)
    produced_assets = _produced_assets(
        graph,
        produced_by_ref=job_config["produced_by_ref"],
        result_contract_order=tuple(job_config["result_contract_order"]),
    )
    inlets = _shared_inlets_for_outputs(graph, produced_assets)

    return PythonAssetJobConfig(
        dag_id=job_config["dag_id"],
        task_id=job_config["task_id"],
        owner=job_config["owner"],
        start_date=job_config["start_date"],
        tags=tuple(job_config["tags"]),
        catchup=bool(job_config.get("catchup", False)),
        max_active_runs=int(job_config.get("max_active_runs", 1)),
        callable_path=job_config["callable_path"],
        runner_executable_path=_resolve_runtime_path(runtime_root, job_config["runner_executable_path"]),
        produced_by_ref=job_config["produced_by_ref"],
        run_kwargs=dict(job_config.get("run_kwargs", {})),
        inlets=tuple(inlets),
        outlets=tuple(asset.airflow for asset in produced_assets),
        result_contract_order=tuple(job_config["result_contract_order"]),
    )


def run_python_asset_job(config: PythonAssetJobConfig) -> dict[str, object]:
    payload = json.dumps(
        {
            "callable_path": config.callable_path,
            "run_kwargs": config.run_kwargs,
        },
        sort_keys=True,
    )
    process = subprocess.Popen(
        [str(config.runner_executable_path), "--payload-json", payload],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
    )
    assert process.stdout is not None

    output_lines: list[str] = []
    for line in process.stdout:
        print(line, end="", flush=True)
        output_lines.append(line)

    returncode = process.wait()
    stdout = "".join(output_lines)
    if returncode != 0:
        raise RuntimeError(
            f"python asset runner failed with exit code {returncode}: {config.runner_executable_path}"
        )

    try:
        row_counts = _row_counts_from_runner_stdout(stdout)
    except PythonAssetJobSkipped as exc:
        from airflow.sdk.exceptions import AirflowSkipException

        raise AirflowSkipException(str(exc)) from exc

    return output_row_count_metadata_from_counts(output_assets=config.outlets, row_counts=row_counts)


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 import_callable(callable_path: str) -> Callable[..., Any]:
    module_name, separator, attr_name = callable_path.partition(":")
    if not separator or not module_name or not attr_name:
        raise ValueError(f"callable path must be '<module>:<attribute>', got {callable_path!r}")
    module = importlib.import_module(module_name)
    value = getattr(module, attr_name)
    if not callable(value):
        raise TypeError(f"{callable_path!r} resolved to a non-callable object")
    return value


def output_row_count_metadata(
    *,
    output_assets: tuple[AirflowAssetSpec, ...],
    results: Any,
) -> dict[str, object]:
    if not isinstance(results, tuple):
        results = (results,)
    if len(results) != len(output_assets):
        raise ValueError(
            f"python asset job returned {len(results)} result(s) for {len(output_assets)} output asset(s)"
        )

    return output_row_count_metadata_from_counts(
        output_assets=output_assets,
        row_counts=tuple(len(result) if hasattr(result, "__len__") else None for result in results),
    )


def output_row_count_metadata_from_counts(
    *,
    output_assets: tuple[AirflowAssetSpec, ...],
    row_counts: tuple[int | None, ...],
) -> dict[str, object]:
    if len(row_counts) != len(output_assets):
        raise ValueError(
            f"python asset runner returned {len(row_counts)} row count(s) for {len(output_assets)} output asset(s)"
        )

    outputs = []
    for asset, row_count in zip(output_assets, row_counts, strict=True):
        outputs.append(
            {
                "asset_name": asset.name,
                "asset_uri": asset.uri,
                "row_count": row_count,
            }
        )
    return {"outputs": outputs}


def attach_output_row_count_metadata(
    *,
    outlet_events: Any,
    airflow_assets: tuple[Any, ...] | list[Any],
    metadata: dict[str, object],
) -> None:
    outputs = metadata.get("outputs")
    if not isinstance(outputs, list):
        raise ValueError("python asset metadata is missing outputs")
    if len(outputs) != len(airflow_assets):
        raise ValueError(
            f"python asset metadata has {len(outputs)} output(s) for {len(airflow_assets)} Airflow outlet(s)"
        )

    for airflow_asset, output in zip(airflow_assets, outputs, strict=True):
        if not isinstance(output, dict):
            raise ValueError("python asset metadata output must be a mapping")
        outlet_events[airflow_asset].extra["row_count"] = output.get("row_count")


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 _row_counts_from_runner_stdout(stdout: str) -> tuple[int | None, ...]:
    for line in reversed(stdout.splitlines()):
        try:
            payload = json.loads(line)
        except json.JSONDecodeError:
            continue
        row_counts = payload.get("row_counts")
        if isinstance(row_counts, list):
            return tuple(row_counts)
        skip_reason = payload.get("skip_reason")
        if isinstance(skip_reason, str) and skip_reason:
            raise PythonAssetJobSkipped(skip_reason)
    raise ValueError("python asset runner did not emit a row_counts JSON line")


def _produced_assets(
    graph: CompiledAssetGraph,
    *,
    produced_by_ref: str,
    result_contract_order: tuple[str, ...],
) -> tuple[Any, ...]:
    by_contract_id = {
        asset.contract_id: asset
        for asset in graph.assets.values()
        if asset.produced_by_kind == "python" and asset.produced_by_ref == produced_by_ref
    }
    missing = sorted(contract_id for contract_id in result_contract_order if contract_id not in by_contract_id)
    if missing:
        raise ValueError(
            f"python job {produced_by_ref!r} is missing contract output(s): {', '.join(missing)}"
        )
    extras = sorted(set(by_contract_id) - set(result_contract_order))
    if extras:
        raise ValueError(
            f"python job {produced_by_ref!r} has output contract(s) not listed in config: {', '.join(extras)}"
        )
    return tuple(by_contract_id[contract_id] for contract_id in result_contract_order)


def _shared_inlets_for_outputs(
    graph: CompiledAssetGraph,
    produced_assets: tuple[Any, ...],
) -> list[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 [inlet_by_name[name] for name in sorted(inlet_by_name)]