from __future__ import annotations

from pathlib import Path

import yaml


REPO_ROOT = Path(__file__).resolve().parents[1]
AIRFLOW_ROOT = REPO_ROOT / "orchestration" / "airflow"


def test_airflow_local_compose_declares_required_services() -> None:
    compose = yaml.safe_load((AIRFLOW_ROOT / "setup" / "docker-compose.yaml").read_text())
    services = compose["services"]
    common = compose["x-airflow-common"]

    assert common["image"] == "${AIRFLOW_IMAGE_NAME:-pingcap-datalake-airflow:3.2.1}"
    assert common["build"]["dockerfile"] == "Dockerfile"
    assert common["build"]["args"]["AIRFLOW_BASE_IMAGE"] == (
        "${AIRFLOW_BASE_IMAGE:-apache/airflow:3.2.1}"
    )
    for service in (
        "postgres",
        "redis",
        "airflow-apiserver",
        "airflow-scheduler",
        "airflow-dag-processor",
        "airflow-worker",
        "airflow-triggerer",
        "airflow-init",
    ):
        assert service in services

    assert services["airflow-triggerer"]["command"] == "triggerer"
    assert services["airflow-apiserver"]["command"] == "api-server"
    assert services["airflow-scheduler"]["command"] == "scheduler"
    assert services["airflow-dag-processor"]["command"] == "dag-processor"
    assert services["airflow-worker"]["command"] == "celery worker"
    assert common["environment"]["UV_CACHE_DIR"] == "/opt/airflow/.cache/uv"
    assert common["environment"]["DATALAKE_AIRFLOW_REPO_ROOT"] == (
        "${DATALAKE_AIRFLOW_REPO_ROOT:-/opt/datalake}"
    )
    assert common["environment"]["DATALAKE_LOAD_WORKSPACE"] == (
        "${DATALAKE_LOAD_WORKSPACE:-/opt/datalake-runtime/workspace}"
    )
    assert common["environment"]["DATALAKE_EXTRACT_WORKSPACE"] == (
        "${DATALAKE_EXTRACT_WORKSPACE:-/opt/datalake-runtime/extract-workspace}"
    )
    assert "UV_PROJECT_ENVIRONMENT" not in common["environment"]
    assert "airflow-uv-cache:/opt/airflow/.cache/uv" in common["volumes"]
    assert "airflow-venvs:/opt/airflow/.venvs" in common["volumes"]
    assert "datalake-runtime:/opt/datalake-runtime" in common["volumes"]
    assert "airflow-uv-cache" in compose["volumes"]
    assert "airflow-venvs" in compose["volumes"]
    assert "datalake-runtime" in compose["volumes"]


def test_airflow_layout_has_engine_config_dag_and_setup_dirs() -> None:
    for directory in ("engine", "dags", "sources", "setup", "logs", "plugins"):
        assert (AIRFLOW_ROOT / directory).is_dir()

    assert (REPO_ROOT / "orchestration" / "README.md").is_file()


def test_hello_assets_dag_uses_asset_aware_scheduling() -> None:
    dag_source = (AIRFLOW_ROOT / "dags" / "hello_assets.py").read_text()
    source_config = yaml.safe_load((AIRFLOW_ROOT / "sources" / "local-dev.yaml").read_text())

    assert "from airflow.sdk import DAG, task" in dag_source
    assert "load_hello_asset_smoke_config" in dag_source
    assert "x-datalake://orchestration/hello_raw" not in dag_source
    assert "x-datalake://orchestration/hello_consumed" not in dag_source
    assert "schedule=HELLO_ASSETS.producer_schedule" in dag_source
    assert "schedule=HELLO_ASSETS.consumer_schedule" in dag_source
    assert "@task(outlets=[HELLO_ASSETS.raw_asset])" in dag_source
    assert (
        "@task(inlets=[HELLO_ASSETS.raw_asset], outlets=[HELLO_ASSETS.consumed_asset])"
        in dag_source
    )
    assert (
        source_config["hello_assets"]["raw_asset"]["uri"]
        == "x-datalake://orchestration/hello_raw"
    )
    assert (
        source_config["hello_assets"]["consumed_asset"]["uri"]
        == "x-datalake://orchestration/hello_consumed"
    )