from __future__ import annotations

import json
from pathlib import Path
import sys

import yaml


REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT))

from orchestration.airflow.engine.contract_assets import AirflowAssetSpec
from orchestration.airflow.engine.load_assets import (
    build_airflow_asset,
    execute_load_asset_command,
    load_load_asset_config,
    resolve_manifest_path,
)


LOAD_ASSETS_CONFIG_PATH = REPO_ROOT / "orchestration" / "airflow" / "sources" / "load-assets.yaml"


def test_manifest_load_unit_binds_source_to_raw_contract_assets() -> None:
    config = load_load_asset_config(LOAD_ASSETS_CONFIG_PATH, repo_root=REPO_ROOT)
    unit = config.by_name()["data-foundation-finance-manifest-load"]

    assert unit.mode == "manifest"
    assert unit.schedule[0].name == "ingestion.extract.tidbcloud.data_foundation_finance.manifest"
    assert unit.schedule[0].group == "manifest"
    assert unit.manifest_path_env == "DATA_FOUNDATION_FINANCE_LOAD_MANIFEST"
    assert unit.manifest_event_extra_keys == ("manifest_uri", "manifest_path")
    assert "source.data_foundation_finance.salesforce_order_s3_snapshot" in {
        asset.name for asset in unit.inlets
    }
    assert "finance.data_foundation_finance.salesforce_order_snapshot" in {
        asset.name for asset in unit.outlets
    }
    assert all(asset.group == "raw" for asset in unit.outlets)
    assert len(unit.outlets) == 27

    argv = unit.command.argv(manifest_path="/tmp/manifest.json")
    assert argv[:5] == [
        str(REPO_ROOT / "orchestration" / "airflow" / "bin" / "load-runner"),
        "ingestion.load.s3_bulk_load.engine.cli",
        "--source",
        "ingestion/load/s3_bulk_load/sources/data_foundation_finance/source.yaml",
        "load-run",
    ]
    assert argv[-2:] == ["--manifest", "/tmp/manifest.json"]


def test_nexuscrm_manifest_load_unit_binds_source_to_raw_contract_assets() -> None:
    config = load_load_asset_config(LOAD_ASSETS_CONFIG_PATH, repo_root=REPO_ROOT)
    unit = config.by_name()["nexuscrm-manifest-load"]

    assert unit.mode == "manifest"
    assert unit.dag_id == "orchestration_load_nexuscrm"
    assert unit.task_id == "load_nexuscrm_manifest"
    assert unit.schedule[0].name == "ingestion.extract.tidbcloud.nexuscrm.manifest"
    assert unit.schedule[0].uri == "x-datalake://ingestion/manifest/tidbcloud/nexuscrm"
    assert unit.schedule[0].group == "manifest"
    assert unit.manifest_path_env == "NEXUSCRM_LOAD_MANIFEST"
    assert unit.manifest_event_extra_keys == ("manifest_uri", "manifest_path")
    assert "source.nexuscrm.account_s3_snapshot" in {asset.name for asset in unit.inlets}
    assert "crm.nexuscrm.account_snapshot" in {asset.name for asset in unit.outlets}
    assert all(asset.group == "raw" for asset in unit.outlets)
    assert len(unit.outlets) == 196

    argv = unit.command.argv(manifest_path="/tmp/manifest.json")
    assert argv[:5] == [
        str(REPO_ROOT / "orchestration" / "airflow" / "bin" / "load-runner"),
        "ingestion.load.s3_bulk_load.engine.cli",
        "--source",
        "ingestion/load/s3_bulk_load/sources/nexuscrm/source.yaml",
        "load-run",
    ]
    assert argv[-2:] == ["--manifest", "/tmp/manifest.json"]


def test_external_incremental_load_unit_uses_poll_schedule_and_contract_assets() -> None:
    config = load_load_asset_config(LOAD_ASSETS_CONFIG_PATH, repo_root=REPO_ROOT)
    unit = config.by_name()["metadb-external-incremental-poll"]

    assert unit.mode == "external_incremental"
    assert unit.schedule == "17 * * * *"
    assert [asset.name for asset in unit.inlets] == ["source.metadb.billing2_invoices_external"]
    assert [asset.name for asset in unit.outlets] == ["billing.metadb.billing2_invoices_snapshot"]
    assert unit.outlets[0].uri == "x-datalake://billing/metadb/billing2_invoices_snapshot"
    assert unit.outlets[0].group == "raw"
    assert unit.command.argv() == [
        str(REPO_ROOT / "orchestration" / "airflow" / "bin" / "load-runner"),
        "ingestion.load.gcs_bulk_load.engine.cli",
        "--source",
        "ingestion/load/gcs_bulk_load/sources/metadb/source.yaml",
        "load-run",
    ]


def test_manifest_path_resolves_from_asset_event_metadata_and_env() -> None:
    config = load_load_asset_config(LOAD_ASSETS_CONFIG_PATH, repo_root=REPO_ROOT)
    unit = config.by_name()["data-foundation-finance-manifest-load"]

    context_path = resolve_manifest_path(
        unit,
        context={
            "triggering_asset_events": {
                unit.manifest_asset.uri: [
                    {
                        "extra": {
                            "manifest_path": (
                                "ingestion/load/s3_bulk_load/sources/"
                                "data_foundation_finance/runs/run-1/manifest.json"
                            )
                        }
                    }
                ]
            }
        },
        env={},
    )
    env_path = resolve_manifest_path(
        unit,
        context={},
        env={"DATA_FOUNDATION_FINANCE_LOAD_MANIFEST": "/tmp/env-manifest.json"},
    )

    assert context_path == (
        REPO_ROOT / "ingestion/load/s3_bulk_load/sources/data_foundation_finance/runs/run-1/manifest.json"
    )
    assert env_path == Path("/tmp/env-manifest.json")


def test_manifest_event_prefers_remote_manifest_uri(monkeypatch) -> None:
    config = load_load_asset_config(LOAD_ASSETS_CONFIG_PATH, repo_root=REPO_ROOT)
    unit = config.by_name()["data-foundation-finance-manifest-load"]
    downloads = []

    def fake_download(uri: str, destination: Path) -> None:
        downloads.append((uri, destination))
        destination.parent.mkdir(parents=True, exist_ok=True)
        destination.write_text("{}")

    monkeypatch.setattr("orchestration.airflow.engine.load_assets._download_object", fake_download)

    path = resolve_manifest_path(
        unit,
        context={
            "triggering_asset_events": {
                unit.manifest_asset.uri: [
                    {
                        "extra": {
                            "manifest_uri": "s3://bucket/path/manifest.json",
                            "manifest_path": (
                                "/opt/datalake-runtime/extract-workspace/"
                                "runs/run-1/manifest.json"
                            ),
                        }
                    }
                ]
            }
        },
        env={},
    )

    assert downloads[0][0] == "s3://bucket/path/manifest.json"
    assert path == downloads[0][1]
    assert path.read_text() == "{}"


def test_relative_manifest_paths_resolve_against_load_workspace(monkeypatch) -> None:
    monkeypatch.setenv("DATALAKE_AIRFLOW_REPO_ROOT", "/opt/datalake")
    monkeypatch.setenv("DATALAKE_LOAD_WORKSPACE", "/opt/datalake-runtime/workspace")
    config = load_load_asset_config(LOAD_ASSETS_CONFIG_PATH, repo_root=REPO_ROOT)
    unit = config.by_name()["data-foundation-finance-manifest-load"]

    path = resolve_manifest_path(
        unit,
        context={},
        env={"DATA_FOUNDATION_FINANCE_LOAD_MANIFEST": "runs/run-1/manifest.json"},
    )

    assert unit.runtime_repo_root == Path("/opt/datalake")
    assert unit.runtime_workspace_root == Path("/opt/datalake-runtime/workspace")
    assert unit.command.argv(manifest_path=path)[3] == (
        "ingestion/load/s3_bulk_load/sources/data_foundation_finance/source.yaml"
    )
    assert path == Path("/opt/datalake-runtime/workspace/runs/run-1/manifest.json")

    source_root_path = (
        "/opt/datalake/ingestion/load/s3_bulk_load/sources/"
        "data_foundation_finance/runs/run-1/manifest.json"
    )
    remapped_path = resolve_manifest_path(
        unit,
        context={},
        env={"DATA_FOUNDATION_FINANCE_LOAD_MANIFEST": source_root_path},
    )

    assert remapped_path == Path(
        "/opt/datalake-runtime/workspace/ingestion/load/s3_bulk_load/sources/"
        "data_foundation_finance/runs/run-1/manifest.json"
    )


def test_configured_manifest_path_resolves_against_load_workspace(monkeypatch, tmp_path) -> None:
    monkeypatch.setenv("DATALAKE_AIRFLOW_REPO_ROOT", "/opt/datalake")
    monkeypatch.setenv("DATALAKE_LOAD_WORKSPACE", "/opt/datalake-runtime/workspace")
    config_data = yaml.safe_load(LOAD_ASSETS_CONFIG_PATH.read_text())
    config_data["load_asset_units"][0]["manifest"]["path"] = "runs/configured/manifest.json"
    config_data["load_asset_units"][0]["manifest"].pop("path_env")
    config_path = tmp_path / "load-assets.yaml"
    config_path.write_text(yaml.safe_dump(config_data))

    config = load_load_asset_config(config_path, repo_root=REPO_ROOT)
    unit = config.by_name()["data-foundation-finance-manifest-load"]

    assert unit.manifest_path == Path("/opt/datalake-runtime/workspace/runs/configured/manifest.json")
    assert resolve_manifest_path(unit, context={}, env={}) == unit.manifest_path


def test_execute_load_asset_command_returns_output_asset_metadata(monkeypatch) -> None:
    config = load_load_asset_config(LOAD_ASSETS_CONFIG_PATH, repo_root=REPO_ROOT)
    unit = config.by_name()["data-foundation-finance-manifest-load"]

    def fake_run(argv, **kwargs):
        assert argv[-2:] == ["--manifest", "/tmp/manifest.json"]

        class Completed:
            stdout = (
                '{"loads": [{"table": "salesforce_order", '
                '"raw_table": "data_foundation_finance_raw.data_foundation_finance_salesforce_order_snapshots", '
                '"loaded_rows": 10}]}'
            )

        return Completed()

    monkeypatch.setattr("orchestration.airflow.engine.load_assets.subprocess.run", fake_run)

    metadata = execute_load_asset_command(
        unit.command,
        output_assets=(
            AirflowAssetSpec(
                name="finance.data_foundation_finance.salesforce_order_snapshot",
                uri="x-datalake://finance/data_foundation_finance/salesforce_order_snapshot",
                group="raw",
                extra={"asset_ref": "data_foundation_finance_raw.data_foundation_finance_salesforce_order_snapshots"},
            ),
        ),
        manifest_path="/tmp/manifest.json",
    )

    assert metadata == {
        "outputs": [
            {
                "asset_name": "finance.data_foundation_finance.salesforce_order_snapshot",
                "asset_uri": "x-datalake://finance/data_foundation_finance/salesforce_order_snapshot",
                "asset_group": "raw",
            }
        ],
        "loader": {
            "loads": [
                {
                    "table": "salesforce_order",
                    "raw_table": (
                        "data_foundation_finance_raw."
                        "data_foundation_finance_salesforce_order_snapshots"
                    ),
                    "loaded_rows": 10,
                }
            ]
        },
    }


def test_manifest_load_command_fails_when_a_declared_outlet_is_not_loaded(monkeypatch) -> None:
    config = load_load_asset_config(LOAD_ASSETS_CONFIG_PATH, repo_root=REPO_ROOT)
    unit = config.by_name()["data-foundation-finance-manifest-load"]

    def fake_run(argv, **kwargs):
        class Completed:
            stdout = '{"loads": []}'

        return Completed()

    monkeypatch.setattr("orchestration.airflow.engine.load_assets.subprocess.run", fake_run)

    try:
        execute_load_asset_command(
            unit.command,
            output_assets=(
                AirflowAssetSpec(
                    name="finance.data_foundation_finance.salesforce_order_snapshot",
                    uri="x-datalake://finance/data_foundation_finance/salesforce_order_snapshot",
                    group="raw",
                    extra={
                        "asset_ref": (
                            "data_foundation_finance_raw."
                            "data_foundation_finance_salesforce_order_snapshots"
                        )
                    },
                ),
            ),
            manifest_path="/tmp/manifest.json",
        )
    except RuntimeError as error:
        assert "did not materialize every declared outlet" in str(error)
    else:
        raise AssertionError("manifest load should fail when a declared outlet is not loaded")


def test_external_incremental_noop_poll_skips_without_emitting_assets(monkeypatch) -> None:
    import pytest

    airflow_exceptions = pytest.importorskip("airflow.sdk.exceptions")
    config = load_load_asset_config(LOAD_ASSETS_CONFIG_PATH, repo_root=REPO_ROOT)
    unit = config.by_name()["metadb-external-incremental-poll"]

    def fake_run(argv, **kwargs):
        class Completed:
            stdout = '{"loads": []}'

        return Completed()

    monkeypatch.setattr("orchestration.airflow.engine.load_assets.subprocess.run", fake_run)

    with pytest.raises(airflow_exceptions.AirflowSkipException):
        execute_load_asset_command(
            unit.command,
            output_assets=unit.outlets,
        )


def test_failed_load_command_includes_captured_output_and_pr_request(
    monkeypatch,
    tmp_path,
) -> None:
    monkeypatch.setenv("DATALAKE_LOAD_WORKSPACE", str(tmp_path))
    config = load_load_asset_config(LOAD_ASSETS_CONFIG_PATH, repo_root=REPO_ROOT)
    unit = config.by_name()["metadb-external-incremental-poll"]

    def fake_run(argv, **kwargs):
        raise __import__("subprocess").CalledProcessError(
            returncode=3,
            cmd=argv,
            output="",
            stderr="schema drift wrote contracts/proposed/metadb.contract.yaml",
        )

    monkeypatch.setattr("orchestration.airflow.engine.load_assets.subprocess.run", fake_run)

    try:
        execute_load_asset_command(unit.command, output_assets=unit.outlets)
    except RuntimeError as error:
        assert "exit code 3" in str(error)
        assert "contracts/proposed/metadb.contract.yaml" in str(error)
        assert "contract-update PR action request:" in str(error)
    else:
        raise AssertionError("failed load command should surface captured stderr")

    action_requests = list(
        (
            tmp_path
            / "ingestion/load/gcs_bulk_load/sources/metadb/runs/action-requests"
        ).glob("*.json")
    )
    assert len(action_requests) == 1
    request = json.loads(action_requests[0].read_text())
    assert request["action"] == "open_pr"
    assert request["risk_tier"] == "controlled_operational"
    assert request["parameters"]["source_name"] == "metadb"
    assert request["parameters"]["proposed_contract_path"] == (
        "contracts/proposed/metadb.contract.yaml"
    )
    assert request["parameters"]["promote_to_path"] == (
        "contracts/upstream/metadb.contract.yaml"
    )
    assert request["parameters"]["branch_name"] == (
        "codex/contract-update-metadb-schema-drift"
    )


def test_manifest_download_failure_includes_captured_output(monkeypatch) -> None:
    config = load_load_asset_config(LOAD_ASSETS_CONFIG_PATH, repo_root=REPO_ROOT)
    unit = config.by_name()["data-foundation-finance-manifest-load"]

    def fake_run(argv, **kwargs):
        raise __import__("subprocess").CalledProcessError(
            returncode=4,
            cmd=argv,
            output="",
            stderr="permission denied for manifest",
        )

    monkeypatch.setattr("orchestration.airflow.engine.load_assets.subprocess.run", fake_run)

    try:
        resolve_manifest_path(
            unit,
            context={
                "triggering_asset_events": {
                    unit.manifest_asset.uri: [
                        {"extra": {"manifest_uri": "gs://bucket/path/manifest.json"}}
                    ]
                }
            },
            env={},
        )
    except RuntimeError as error:
        assert "gcloud storage cp gs://bucket/path/manifest.json" in str(error)
        assert "permission denied for manifest" in str(error)
    else:
        raise AssertionError("manifest download failures should surface captured stderr")


def test_load_assets_dag_is_thin_contract_asset_registration() -> None:
    dag_source = (
        REPO_ROOT / "orchestration" / "airflow" / "dags" / "ingestion_load_assets.py"
    ).read_text()
    engine_source = (
        REPO_ROOT / "orchestration" / "airflow" / "engine" / "load_assets.py"
    ).read_text()
    runner_source = (REPO_ROOT / "orchestration" / "airflow" / "bin" / "load-runner").read_text()

    assert "load_load_asset_config(LOAD_ASSETS_CONFIG_PATH)" in dag_source
    assert "REPO_ROOT / \"orchestration\" / \"airflow\" / \"sources\" / \"load-assets.yaml\"" in dag_source
    assert "resolve_manifest_path(unit, context=get_current_context())" in dag_source
    assert "execute_load_asset_command(" in dag_source
    assert "finance.data_foundation_finance.salesforce_order_snapshot" not in dag_source
    assert "billing.metadb.billing2_invoices_snapshot" not in dag_source
    assert "data_foundation_finance" not in engine_source
    assert "metadb" not in engine_source
    assert "python -m" in runner_source
    assert "[ -f \"${DATALAKE_ROOT}/.env\" ]" in runner_source
    assert "cd \"${DATALAKE_ROOT}\"" in runner_source
    assert "--exclude \"*/contracts/proposed/***\"" in runner_source
    assert "--exclude \"*/runs/***\"" in runner_source


def test_build_load_airflow_asset_preserves_contract_group() -> None:
    import pytest

    pytest.importorskip("airflow.sdk")

    asset = build_airflow_asset(
        AirflowAssetSpec(
            name="billing.metadb.billing2_invoices_snapshot",
            uri="x-datalake://billing/metadb/billing2_invoices_snapshot",
            group="raw",
            extra={"contract_id": "billing.metadb.billing2_invoices_snapshot"},
        )
    )

    assert asset.group == "raw"