from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

import yaml


CONTRACT_GLOBS = (
    "ingestion/**/contracts/**/*.contract.yaml",
    "ingestion/**/contracts/**/*.contract.yml",
    "transform/**/contracts/**/*.contract.yaml",
    "transform/**/contracts/**/*.contract.yml",
)


@dataclass(frozen=True)
class AirflowAssetSpec:
    name: str
    uri: str
    group: str
    extra: dict[str, str]

    def to_kwargs(self) -> dict[str, object]:
        return {
            "name": self.name,
            "uri": self.uri,
            "group": self.group,
            "extra": dict(self.extra),
        }


@dataclass(frozen=True)
class CompiledContractAsset:
    contract_id: str
    contract_path: str
    contract_version: str
    contract_kind: str
    domain: str
    layer: str | None
    asset_id: str
    asset_kind: str
    asset_ref: str
    materialization: str | None
    produced_by_kind: str
    produced_by_ref: str
    freshness: dict[str, str]
    airflow: AirflowAssetSpec


@dataclass(frozen=True)
class CompiledAssetEdge:
    upstream_ref: str
    upstream_kind: str
    downstream_ref: str
    required: bool
    upstream_asset_name: str | None
    downstream_asset_name: str


@dataclass(frozen=True)
class ExternalUpstreamAsset:
    ref: str
    upstream_kind: str
    airflow: AirflowAssetSpec


@dataclass(frozen=True)
class CompiledAssetGraph:
    assets: dict[str, CompiledContractAsset]
    external_assets: dict[str, ExternalUpstreamAsset]
    edges: tuple[CompiledAssetEdge, ...]

    def asset_names(self) -> set[str]:
        return {
            *(asset.airflow.name for asset in self.assets.values()),
            *(asset.airflow.name for asset in self.external_assets.values()),
        }

    def to_manifest(self) -> dict[str, object]:
        return {
            "assets": [
                {
                    "contract_id": asset.contract_id,
                    "contract_path": asset.contract_path,
                    "contract_version": asset.contract_version,
                    "contract_kind": asset.contract_kind,
                    "domain": asset.domain,
                    "layer": asset.layer,
                    "asset_id": asset.asset_id,
                    "asset_kind": asset.asset_kind,
                    "asset_ref": asset.asset_ref,
                    "materialization": asset.materialization,
                    "produced_by": {
                        "kind": asset.produced_by_kind,
                        "ref": asset.produced_by_ref,
                    },
                    "freshness": dict(asset.freshness),
                    "airflow": asset.airflow.to_kwargs(),
                }
                for asset in sorted(self.assets.values(), key=lambda item: item.contract_id)
            ],
            "external_assets": [
                {
                    "ref": asset.ref,
                    "upstream_kind": asset.upstream_kind,
                    "airflow": asset.airflow.to_kwargs(),
                }
                for asset in sorted(self.external_assets.values(), key=lambda item: item.ref)
            ],
            "edges": [
                {
                    "upstream_ref": edge.upstream_ref,
                    "upstream_kind": edge.upstream_kind,
                    "downstream_ref": edge.downstream_ref,
                    "required": edge.required,
                    "upstream_asset_name": edge.upstream_asset_name,
                    "downstream_asset_name": edge.downstream_asset_name,
                }
                for edge in sorted(
                    self.edges,
                    key=lambda item: (
                        item.downstream_ref,
                        item.upstream_kind,
                        item.upstream_ref,
                    ),
                )
            ],
        }


def discover_contract_paths(
    repo_root: Path,
    contract_globs: tuple[str, ...] = CONTRACT_GLOBS,
) -> tuple[Path, ...]:
    paths: set[Path] = set()
    for pattern in contract_globs:
        paths.update(path for path in repo_root.glob(pattern) if path.is_file())
    return tuple(sorted(paths))


def compile_repo_contract_assets(repo_root: Path) -> CompiledAssetGraph:
    return compile_contract_assets(
        discover_contract_paths(repo_root),
        repo_root=repo_root,
    )


def compile_contract_assets(
    contract_paths: tuple[Path, ...] | list[Path],
    *,
    repo_root: Path,
) -> CompiledAssetGraph:
    compiled_assets: dict[str, CompiledContractAsset] = {}
    declared_upstreams: list[tuple[str, dict[str, object]]] = []

    for path in sorted(contract_paths):
        contract = _load_contract(path)
        asset = _compile_contract_asset(contract, path, repo_root)
        if asset.contract_id in compiled_assets:
            raise ValueError(f"Duplicate contract id {asset.contract_id!r} at {path}")
        compiled_assets[asset.contract_id] = asset
        for upstream in contract.get("declared_lineage", {}).get("upstreams", []):
            declared_upstreams.append((asset.contract_id, upstream))

    external_assets = _external_assets_for_upstreams(declared_upstreams, compiled_assets)
    edges = tuple(
        _compile_edge(downstream_ref, upstream, compiled_assets, external_assets)
        for downstream_ref, upstream in declared_upstreams
    )

    return CompiledAssetGraph(
        assets=compiled_assets,
        external_assets=external_assets,
        edges=edges,
    )


def stale_asset_names(
    previous_asset_names: set[str] | list[str] | tuple[str, ...],
    current_graph: CompiledAssetGraph,
) -> set[str]:
    return set(previous_asset_names) - current_graph.asset_names()


def _load_contract(path: Path) -> dict[str, Any]:
    loaded = yaml.safe_load(path.read_text()) or {}
    if not isinstance(loaded, dict):
        raise ValueError(f"Contract {path} did not parse to a mapping")
    return loaded


def _compile_contract_asset(
    contract: dict[str, Any],
    path: Path,
    repo_root: Path,
) -> CompiledContractAsset:
    contract_id = _required_str(contract, "id", path)
    contract_kind = str(contract.get("kind", "product"))
    domain = _required_nested_str(contract, ("domain", "name"), path)
    layer = contract.get("layer")
    assets = contract.get("assets") or []
    if len(assets) != 1:
        raise ValueError(
            f"Contract {contract_id!r} must declare exactly one asset for the "
            f"v0 orchestration compilation profile; found {len(assets)} in {path}"
        )

    declared_asset = assets[0]
    produced_by = _required_profile_mapping(declared_asset, "produced_by", path)
    asset_group = str(layer) if contract_kind == "product" else "source"
    airflow = AirflowAssetSpec(
        name=contract_id,
        uri=_asset_uri(domain, contract_id),
        group=asset_group,
        extra=_asset_extra(contract, declared_asset, path, repo_root),
    )

    return CompiledContractAsset(
        contract_id=contract_id,
        contract_path=_relative_path(path, repo_root),
        contract_version=_required_str(contract, "version", path),
        contract_kind=contract_kind,
        domain=domain,
        layer=str(layer) if layer is not None else None,
        asset_id=_required_str(declared_asset, "id", path),
        asset_kind=_required_str(declared_asset, "kind", path),
        asset_ref=_required_str(declared_asset, "ref", path),
        materialization=declared_asset.get("materialization"),
        produced_by_kind=_required_str(produced_by, "kind", path),
        produced_by_ref=_required_str(produced_by, "ref", path),
        freshness={
            key: str(value)
            for key, value in (contract.get("freshness") or {}).items()
            if value is not None
        },
        airflow=airflow,
    )


def _compile_edge(
    downstream_ref: str,
    upstream: dict[str, object],
    compiled_assets: dict[str, CompiledContractAsset],
    external_assets: dict[str, ExternalUpstreamAsset],
) -> CompiledAssetEdge:
    upstream_ref = str(upstream["ref"])
    upstream_asset = compiled_assets.get(upstream_ref)
    external_asset = external_assets.get(upstream_ref)
    downstream_asset = compiled_assets[downstream_ref]

    return CompiledAssetEdge(
        upstream_ref=upstream_ref,
        upstream_kind=str(upstream["kind"]),
        downstream_ref=downstream_ref,
        required=bool(upstream.get("required", True)),
        upstream_asset_name=(
            upstream_asset.airflow.name
            if upstream_asset
            else external_asset.airflow.name
            if external_asset
            else None
        ),
        downstream_asset_name=downstream_asset.airflow.name,
    )


def _external_assets_for_upstreams(
    declared_upstreams: list[tuple[str, dict[str, object]]],
    compiled_assets: dict[str, CompiledContractAsset],
) -> dict[str, ExternalUpstreamAsset]:
    external_assets: dict[str, ExternalUpstreamAsset] = {}
    for _, upstream in declared_upstreams:
        upstream_kind = str(upstream["kind"])
        upstream_ref = str(upstream["ref"])
        if upstream_kind not in {"data_product", "source"}:
            continue
        if upstream_ref in compiled_assets or upstream_ref in external_assets:
            continue
        external_assets[upstream_ref] = ExternalUpstreamAsset(
            ref=upstream_ref,
            upstream_kind=upstream_kind,
            airflow=AirflowAssetSpec(
                name=upstream_ref,
                uri=_external_asset_uri(upstream_kind, upstream_ref),
                group="external",
                extra={
                    "external_ref": upstream_ref,
                    "external_kind": upstream_kind,
                    "resolution": "uncontracted_declared_lineage_upstream",
                },
            ),
        )
    return external_assets


def _asset_extra(
    contract: dict[str, Any],
    declared_asset: dict[str, Any],
    path: Path,
    repo_root: Path,
) -> dict[str, str]:
    produced_by = _required_profile_mapping(declared_asset, "produced_by", path)
    freshness = contract.get("freshness") or {}
    layer = contract.get("layer")
    extra = {
        "contract_id": _required_str(contract, "id", path),
        "contract_version": _required_str(contract, "version", path),
        "contract_kind": str(contract.get("kind", "product")),
        "contract_path": _relative_path(path, repo_root),
        "domain": _required_nested_str(contract, ("domain", "name"), path),
        "asset_id": _required_str(declared_asset, "id", path),
        "asset_kind": _required_str(declared_asset, "kind", path),
        "asset_ref": _required_str(declared_asset, "ref", path),
        "produced_by_kind": _required_str(produced_by, "kind", path),
        "produced_by_ref": _required_str(produced_by, "ref", path),
        "freshness_max_age": str(freshness["max_age"]),
    }
    if layer is not None:
        extra["layer"] = str(layer)
    if declared_asset.get("materialization"):
        extra["materialization"] = str(declared_asset["materialization"])
    for optional_key in ("schedule_ref", "degrade_after", "timezone"):
        if freshness.get(optional_key):
            extra[f"freshness_{optional_key}"] = str(freshness[optional_key])
    return extra


def _asset_uri(domain: str, contract_id: str) -> str:
    if contract_id.startswith(f"{domain}."):
        product_path = contract_id[len(domain) + 1 :]
    else:
        product_path = contract_id
    return f"x-datalake://{domain}/{product_path.replace('.', '/')}"


def _external_asset_uri(upstream_kind: str, upstream_ref: str) -> str:
    return f"x-datalake://external/{upstream_kind}/{upstream_ref.replace('.', '/')}"


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


def _required_str(mapping: dict[str, Any], key: str, path: Path) -> str:
    value = mapping.get(key)
    if not isinstance(value, str) or not value:
        raise ValueError(f"Contract {path} is missing required string {key!r}")
    return value


def _required_profile_mapping(mapping: dict[str, Any], key: str, path: Path) -> dict[str, Any]:
    value = mapping.get(key)
    if not isinstance(value, dict) or not value:
        raise ValueError(
            f"Contract {path} is missing required mapping {key!r} for the "
            "v0 orchestration compilation profile"
        )
    return value


def _required_nested_str(
    mapping: dict[str, Any],
    keys: tuple[str, ...],
    path: Path,
) -> str:
    value: Any = mapping
    for key in keys:
        if not isinstance(value, dict):
            raise ValueError(f"Contract {path} is missing required path {'.'.join(keys)!r}")
        value = value.get(key)
    if not isinstance(value, str) or not value:
        raise ValueError(f"Contract {path} is missing required path {'.'.join(keys)!r}")
    return value