from __future__ import annotations

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

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 DbtManifestNode:
    unique_id: str
    name: str
    resource_type: str
    depends_on: tuple[str, ...]


@dataclass(frozen=True)
class DbtNodeAssetBinding:
    node_unique_id: str
    node_name: str
    resource_type: str
    contract_id: str
    outlet: AirflowAssetSpec
    inlets: tuple[AirflowAssetSpec, ...]


@dataclass(frozen=True)
class DbtAssetConfig:
    dag_id: str
    task_group_id: str
    owner: str
    start_date: str
    tags: tuple[str, ...]
    catchup: bool
    max_active_runs: int
    project_path: Path
    manifest_path: Path
    profiles_yml_path: Path
    profile_name: str
    target_name: str
    dbt_executable_path: Path
    select: tuple[str, ...]
    target_models: tuple[str, ...]
    node_bindings: dict[str, DbtNodeAssetBinding]
    schedule_assets: tuple[AirflowAssetSpec, ...]


def load_dbt_asset_config(
    config_path: Path,
    *,
    repo_root: Path = REPO_ROOT,
    config_key: str = "dbt_asset_slice",
) -> DbtAssetConfig:
    source_config = yaml.safe_load(config_path.read_text()) or {}
    dbt_config = source_config[config_key]
    runtime_root = _runtime_repo_root(repo_root)

    project_path = _resolve_runtime_path(runtime_root, dbt_config["project_path"])
    manifest_path = _resolve_runtime_path(runtime_root, dbt_config["manifest_path"])
    profiles_yml_path = _resolve_runtime_path(runtime_root, dbt_config["profiles_yml_path"])
    dbt_executable_path = _resolve_runtime_path(runtime_root, dbt_config["dbt_executable_path"])

    host_manifest_path = _resolve_runtime_path(repo_root, dbt_config["manifest_path"])
    manifest_nodes = _load_manifest_nodes(host_manifest_path)
    selected_node_names = _selected_node_names(
        manifest_nodes,
        target_models=tuple(dbt_config["target_models"]),
    )

    graph = compile_repo_contract_assets(repo_root)
    node_bindings = _node_asset_bindings(
        graph,
        manifest_nodes,
        selected_node_names=selected_node_names,
    )
    schedule_assets = _schedule_assets_for_selected_bindings(graph, node_bindings.values())

    return DbtAssetConfig(
        dag_id=dbt_config["dag_id"],
        task_group_id=dbt_config["task_group_id"],
        owner=dbt_config["owner"],
        start_date=dbt_config["start_date"],
        tags=tuple(dbt_config["tags"]),
        catchup=bool(dbt_config.get("catchup", False)),
        max_active_runs=int(dbt_config.get("max_active_runs", 1)),
        project_path=project_path,
        manifest_path=manifest_path,
        profiles_yml_path=profiles_yml_path,
        profile_name=dbt_config["profile_name"],
        target_name=dbt_config["target_name"],
        dbt_executable_path=dbt_executable_path,
        select=tuple(dbt_config["select"]),
        target_models=tuple(dbt_config["target_models"]),
        node_bindings=node_bindings,
        schedule_assets=schedule_assets,
    )


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 build_contract_asset_converters(
    node_bindings: dict[str, DbtNodeAssetBinding],
) -> dict[Any, Callable[..., Any]]:
    from cosmos.constants import DbtResourceType

    converter = _contract_asset_node_converter(node_bindings)
    return {
        DbtResourceType.MODEL: converter,
        DbtResourceType.SEED: converter,
    }


def _contract_asset_node_converter(node_bindings: dict[str, DbtNodeAssetBinding]) -> Callable[..., Any]:
    def convert_node(
        *,
        dag: Any,
        task_group: Any,
        execution_mode: Any,
        task_args: dict[str, Any],
        node: Any,
        task_id: str,
        **_: Any,
    ) -> Any:
        from cosmos.constants import DbtResourceType, ExecutionMode
        from cosmos.operators.local import DbtRunLocalOperator, DbtSeedLocalOperator

        if execution_mode != ExecutionMode.LOCAL:
            raise ValueError("contract asset converter currently supports Cosmos LOCAL execution only")

        binding = node_bindings.get(node.name)
        if binding is None:
            raise ValueError(f"dbt node {node.name!r} has no contract-backed Airflow Asset binding")

        operator_args = dict(task_args)
        operator_args.update(
            {
                "task_id": task_id,
                "dag": dag,
                "task_group": task_group,
                "select": _node_selector(node),
                "emit_datasets": False,
                "inlets": [build_airflow_asset(asset) for asset in binding.inlets],
                "outlets": [build_airflow_asset(binding.outlet)],
            }
        )

        if node.resource_type == DbtResourceType.MODEL:
            return DbtRunLocalOperator(**operator_args)
        if node.resource_type == DbtResourceType.SEED:
            return DbtSeedLocalOperator(**operator_args)
        raise ValueError(f"unsupported contract asset dbt node type: {node.resource_type!r}")

    return convert_node


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 _load_manifest_nodes(manifest_path: Path) -> dict[str, DbtManifestNode]:
    manifest = json.loads(manifest_path.read_text())
    nodes: dict[str, DbtManifestNode] = {}
    for unique_id, node in manifest.get("nodes", {}).items():
        resource_type = node.get("resource_type")
        if resource_type not in {"model", "seed", "test"}:
            continue
        depends_on = tuple((node.get("depends_on") or {}).get("nodes") or ())
        nodes[unique_id] = DbtManifestNode(
            unique_id=unique_id,
            name=node["name"],
            resource_type=resource_type,
            depends_on=depends_on,
        )
    return nodes


def _selected_node_names(
    manifest_nodes: dict[str, DbtManifestNode],
    *,
    target_models: tuple[str, ...],
) -> set[str]:
    by_name = {
        node.name: node.unique_id
        for node in manifest_nodes.values()
        if node.resource_type in {"model", "seed"}
    }
    missing = sorted(model for model in target_models if model not in by_name)
    if missing:
        raise ValueError(f"target dbt models are missing from manifest: {', '.join(missing)}")

    selected_ids: set[str] = set()

    def visit(unique_id: str) -> None:
        if unique_id in selected_ids:
            return
        node = manifest_nodes.get(unique_id)
        if node is None:
            return
        selected_ids.add(unique_id)
        for upstream_id in node.depends_on:
            visit(upstream_id)

    for model in target_models:
        visit(by_name[model])

    for node in manifest_nodes.values():
        if node.resource_type == "test" and set(node.depends_on).intersection(selected_ids):
            selected_ids.add(node.unique_id)

    return {manifest_nodes[unique_id].name for unique_id in selected_ids}


def _node_asset_bindings(
    graph: CompiledAssetGraph,
    manifest_nodes: dict[str, DbtManifestNode],
    *,
    selected_node_names: set[str],
) -> dict[str, DbtNodeAssetBinding]:
    manifest_by_name = {node.name: node for node in manifest_nodes.values()}
    bindings: dict[str, DbtNodeAssetBinding] = {}

    for asset in graph.assets.values():
        if asset.produced_by_kind != "dbt":
            continue
        node = manifest_by_name.get(asset.produced_by_ref)
        if node is None or node.name not in selected_node_names:
            continue
        bindings[node.name] = DbtNodeAssetBinding(
            node_unique_id=node.unique_id,
            node_name=node.name,
            resource_type=node.resource_type,
            contract_id=asset.contract_id,
            outlet=asset.airflow,
            inlets=tuple(_inlets_for_contract(graph, asset.contract_id)),
        )

    selected_without_bindings = sorted(
        node.name
        for node in manifest_nodes.values()
        if node.name in selected_node_names
        and node.resource_type in {"model", "seed"}
        and node.name not in bindings
    )
    if selected_without_bindings:
        raise ValueError(
            "selected dbt model/seed nodes must have contract-backed asset bindings: "
            + ", ".join(selected_without_bindings)
        )

    return bindings


def _inlets_for_contract(graph: CompiledAssetGraph, contract_id: str) -> list[AirflowAssetSpec]:
    inlets: list[AirflowAssetSpec] = []
    for edge in graph.edges:
        if edge.downstream_ref != contract_id:
            continue
        upstream_asset = graph.assets.get(edge.upstream_ref)
        if upstream_asset:
            inlets.append(upstream_asset.airflow)
            continue
        external_asset = graph.external_assets.get(edge.upstream_ref)
        if external_asset:
            inlets.append(external_asset.airflow)
    return sorted(inlets, key=lambda asset: asset.name)


def _schedule_assets_for_selected_bindings(
    graph: CompiledAssetGraph,
    bindings: Iterable[DbtNodeAssetBinding],
) -> tuple[AirflowAssetSpec, ...]:
    binding_list = tuple(bindings)
    selected_contracts = {binding.contract_id for binding in binding_list}
    schedule_assets: dict[str, AirflowAssetSpec] = {}
    for binding in binding_list:
        for edge in graph.edges:
            if edge.downstream_ref != binding.contract_id:
                continue
            if edge.upstream_ref in selected_contracts:
                continue
            upstream_asset = graph.assets.get(edge.upstream_ref)
            if upstream_asset:
                schedule_assets[upstream_asset.airflow.name] = upstream_asset.airflow
                continue
            external_asset = graph.external_assets.get(edge.upstream_ref)
            if external_asset:
                schedule_assets[external_asset.airflow.name] = external_asset.airflow

    return tuple(schedule_assets[name] for name in sorted(schedule_assets))


def _node_selector(node: Any) -> str:
    fqn = getattr(node, "fqn", None)
    if fqn:
        return f"fqn:{'.'.join(fqn)}"
    return node.resource_name