from __future__ import annotations

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.extract_assets import (
    ExtractAssetCommand,
    build_airflow_asset,
    execute_extract_command,
    load_extract_asset_config,
    manifest_event_extra,
)


EXTRACT_ASSETS_CONFIG_PATH = REPO_ROOT / "orchestration" / "airflow" / "sources" / "extract-assets.yaml"


def test_extract_asset_config_registers_bigquery_sync_manifest_asset() -> None:
    config = load_extract_asset_config(EXTRACT_ASSETS_CONFIG_PATH, repo_root=REPO_ROOT)
    unit = config.by_name()["cloud-tenants-bigquery-sync"]

    assert unit.mode == "bigquery_sync"
    assert unit.schedule is None
    assert unit.manifest_asset.name == "ingestion.extract.bigquery.cloud_tenants.manifest"
    assert unit.manifest_asset.uri == "x-datalake://ingestion/manifest/bigquery/cloud_tenants"
    assert [asset.name for asset in unit.source_outlets] == [
        "source.cloud_tenants.cloud_tenants_gcs_snapshot"
    ]
    assert unit.command.source_path == Path("ingestion/extract/bigquery/sources/cloud_tenants/source.yaml")
    assert unit.command.argv("export") == [
        str(REPO_ROOT / "orchestration" / "airflow" / "bin" / "extract-runner"),
        "ingestion.extract.bigquery.engine.cli",
        "--source",
        "ingestion/extract/bigquery/sources/cloud_tenants/source.yaml",
        "export",
    ]
    assert unit.command.argv("write-manifest", run_date="2026-06-01", run_id="run-1")[-4:] == [
        "--run-date",
        "2026-06-01",
        "--run-id",
        "run-1",
    ]


def test_extract_asset_config_registers_tidbcloud_deferrable_manifest_asset_and_source_outlets() -> None:
    config = load_extract_asset_config(EXTRACT_ASSETS_CONFIG_PATH, repo_root=REPO_ROOT)
    unit = config.by_name()["data-foundation-finance-tidbcloud-deferrable"]

    assert unit.mode == "tidbcloud_deferrable"
    assert unit.schedule == "0 6 1,3,5 * *"
    assert unit.manifest_asset.name == "ingestion.extract.tidbcloud.data_foundation_finance.manifest"
    assert unit.manifest_asset.uri == "x-datalake://ingestion/manifest/tidbcloud/data_foundation_finance"
    assert len(unit.source_outlets) == 27
    assert "source.data_foundation_finance.salesforce_order_s3_snapshot" in {
        asset.name for asset in unit.source_outlets
    }
    assert unit.command.argv("start-export", run_date="2026-06-01", run_id="run-1")[-4:] == [
        "--run-date",
        "2026-06-01",
        "--run-id",
        "run-1",
    ]
    assert unit.command.argv("export-status", export_id="exp-1")[-2:] == ["--export-id", "exp-1"]
    assert unit.command.argv(
        "write-manifest",
        run_date="2026-06-01",
        run_id="run-1",
        export_id="exp-1",
    )[-6:] == ["--run-date", "2026-06-01", "--run-id", "run-1", "--export-id", "exp-1"]


def test_manifest_event_extra_reads_manifest_payload(tmp_path: Path) -> None:
    manifest_path = tmp_path / "manifest.json"
    manifest_path.write_text(
        """
{
  "manifest_uri": "s3://bucket/path/manifest.json",
  "run_date": "2026-06-01",
  "run_id": "run-1",
  "source_name": "data_foundation_finance"
}
""".strip()
    )

    assert manifest_event_extra(
        manifest_payload={"manifest": str(manifest_path)},
        fallback_payload={"run_date": "fallback-date", "run_id": "fallback-id"},
    ) == {
        "manifest_path": str(manifest_path),
        "manifest_uri": "s3://bucket/path/manifest.json",
        "run_date": "2026-06-01",
        "run_id": "run-1",
        "source_name": "data_foundation_finance",
    }


def test_generated_bigquery_source_contract_watermarks_name_declared_columns() -> None:
    contract_path = (
        REPO_ROOT
        / "ingestion"
        / "extract"
        / "bigquery"
        / "sources"
        / "cloud_tenants"
        / "contracts"
        / "generated"
        / "cloud_tenants.cloud_tenants.gcs-snapshot.source.contract.yaml"
    )
    contract = yaml.safe_load(contract_path.read_text())
    declared_columns = {column["name"] for column in contract["schema"]["columns"]}

    assert set(contract["extensions"]["watermark_columns"]) <= declared_columns


def test_execute_extract_command_absolutizes_relative_manifest_paths(monkeypatch, tmp_path: Path) -> None:
    workspace = tmp_path / "extract-workspace"
    manifest_path = workspace / "runs" / "2026-06-01" / "run-1" / "manifest.json"
    manifest_path.parent.mkdir(parents=True)
    manifest_path.write_text(
        """
{
  "manifest_uri": "s3://bucket/path/manifest.json",
  "run_date": "2026-06-01",
  "run_id": "run-1",
  "source_name": "data_foundation_finance"
}
""".strip()
    )
    command = ExtractAssetCommand(
        executable_path=Path("/bin/extract-runner"),
        extractor_module="ingestion.extract.tidbcloud.engine.cli",
        source_path=Path("ingestion/extract/tidbcloud/sources/data_foundation_finance/source.yaml"),
        mode="tidbcloud_deferrable",
        tables=(),
        env={"DATALAKE_EXTRACT_WORKSPACE": str(workspace)},
        cwd=Path("/opt/datalake"),
    )

    def fake_run(argv, **kwargs):
        class Completed:
            stdout = '{"manifest": "runs/2026-06-01/run-1/manifest.json"}'

        return Completed()

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

    manifest_payload = execute_extract_command(
        command,
        "write-manifest",
        run_date="2026-06-01",
        run_id="run-1",
        export_id="exp-1",
    )
    extra = manifest_event_extra(
        manifest_payload=manifest_payload,
        fallback_payload={"run_date": "fallback-date", "run_id": "fallback-id"},
    )

    assert manifest_payload["manifest"] == str(manifest_path)
    assert extra["manifest_path"] == str(manifest_path)
    assert extra["manifest_uri"] == "s3://bucket/path/manifest.json"


def test_extract_asset_config_rejects_sources_without_static_contract_outlets(tmp_path: Path) -> None:
    source_path = tmp_path / "source.yaml"
    source_path.write_text(
        """
source:
  name: missing_contracts
schedule:
  cron: "0 0 * * *"
""".strip()
    )
    config_data = yaml.safe_load(EXTRACT_ASSETS_CONFIG_PATH.read_text())
    config_data["extract_asset_units"] = [
        {
            **config_data["extract_asset_units"][0],
            "name": "missing-contracts-bigquery-sync",
            "source_path": str(source_path),
        }
    ]
    config_data["extract_asset_units"][0].pop("output_contracts", None)
    config_path = tmp_path / "extract-assets.yaml"
    config_path.write_text(yaml.safe_dump(config_data))

    try:
        load_extract_asset_config(config_path, repo_root=REPO_ROOT)
    except ValueError as error:
        assert "has no generated source contracts" in str(error)
    else:
        raise AssertionError("extract sources must have static source contract outlets")


def test_execute_extract_command_surfaces_stderr(monkeypatch) -> None:
    command = ExtractAssetCommand(
        executable_path=Path("/bin/extract-runner"),
        extractor_module="ingestion.extract.bigquery.engine.cli",
        source_path=Path("ingestion/extract/bigquery/sources/cloud_tenants/source.yaml"),
        mode="bigquery_sync",
        tables=(),
        env={},
        cwd=REPO_ROOT,
    )

    def fake_run(argv, **kwargs):
        raise __import__("subprocess").CalledProcessError(
            returncode=2,
            cmd=argv,
            output="",
            stderr="bq extract failed",
        )

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

    try:
        execute_extract_command(command, "export")
    except RuntimeError as error:
        assert "exit code 2" in str(error)
        assert "bq extract failed" in str(error)
    else:
        raise AssertionError("extract command failures should include captured stderr")


def test_extract_assets_dag_is_thin_contract_asset_registration() -> None:
    dag_source = (
        REPO_ROOT / "orchestration" / "airflow" / "dags" / "ingestion_extract_assets.py"
    ).read_text()
    engine_source = (
        REPO_ROOT / "orchestration" / "airflow" / "engine" / "extract_assets.py"
    ).read_text()
    runner_source = (REPO_ROOT / "orchestration" / "airflow" / "bin" / "extract-runner").read_text()

    assert "load_extract_asset_config(EXTRACT_ASSETS_CONFIG_PATH)" in dag_source
    assert "TiDBCloudExportManifestOperator(" in dag_source
    assert "outlet_events[manifest_asset].extra = extra" in dag_source
    assert "x-datalake://ingestion/manifest/tidbcloud/data_foundation_finance" not in dag_source
    assert "data_foundation_finance" not in engine_source
    assert "cloud_tenants" not in engine_source
    assert "DATALAKE_EXTRACT_WORKSPACE" in runner_source
    assert "--exclude \"*/runs/***\"" in runner_source


def test_build_extract_airflow_asset_preserves_group() -> None:
    import pytest

    pytest.importorskip("airflow.sdk")

    asset = build_airflow_asset(
        AirflowAssetSpec(
            name="ingestion.extract.tidbcloud.data_foundation_finance.manifest",
            uri="x-datalake://ingestion/manifest/tidbcloud/data_foundation_finance",
            group="manifest",
            extra={"handoff": "run_manifest"},
        )
    )

    assert asset.group == "manifest"