from __future__ import annotations

from dataclasses import replace
from datetime import date
import json
import os
from pathlib import Path
import subprocess
import sys
from types import SimpleNamespace

import pytest

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.python_assets import (
    attach_output_row_count_metadata,
    build_airflow_asset,
    load_python_asset_job_config,
    output_row_count_metadata,
    run_python_asset_job,
)


PYTHON_REVREC_CONFIG_PATH = REPO_ROOT / "orchestration" / "airflow" / "sources" / "python-revrec.yaml"


def test_python_revrec_config_binds_job_to_contract_assets() -> None:
    config = load_python_asset_job_config(PYTHON_REVREC_CONFIG_PATH, repo_root=REPO_ROOT)

    assert config.callable_path == "transform.python.jobs.monthly_cloud_billing_revrec:run"
    assert config.runner_executable_path == (
        REPO_ROOT / "orchestration" / "airflow" / "bin" / "python-asset-job-shim"
    )
    assert config.produced_by_ref == "monthly_cloud_billing_revrec"
    assert [asset.name for asset in config.inlets] == [
        "finance.salesforce_order_daily_view",
        "finance.salesforce_org_monthly_revrec",
    ]
    assert [asset.name for asset in config.outlets] == [
        "finance.monthly_revrec_collections",
        "finance.monthly_revrec_orders",
    ]
    assert [asset.uri for asset in config.outlets] == [
        "x-datalake://finance/monthly_revrec_collections",
        "x-datalake://finance/monthly_revrec_orders",
    ]
    assert all(asset.group == "curated" for asset in config.inlets + config.outlets)


def test_python_asset_row_count_metadata_follows_contract_output_order() -> None:
    metadata = output_row_count_metadata(
        output_assets=(
            AirflowAssetSpec(
                name="finance.monthly_revrec_collections",
                uri="x-datalake://finance/monthly_revrec_collections",
                group="curated",
                extra={},
            ),
            AirflowAssetSpec(
                name="finance.monthly_revrec_orders",
                uri="x-datalake://finance/monthly_revrec_orders",
                group="curated",
                extra={},
            ),
        ),
        results=([1, 2], [1, 2, 3]),
    )

    assert metadata == {
        "outputs": [
            {
                "asset_name": "finance.monthly_revrec_collections",
                "asset_uri": "x-datalake://finance/monthly_revrec_collections",
                "row_count": 2,
            },
            {
                "asset_name": "finance.monthly_revrec_orders",
                "asset_uri": "x-datalake://finance/monthly_revrec_orders",
                "row_count": 3,
            },
        ]
    }


def test_python_asset_metadata_attaches_row_counts_to_outlet_events() -> None:
    class FakeOutletEvents(dict):
        def __getitem__(self, key):
            if key not in self:
                self[key] = SimpleNamespace(extra={})
            return super().__getitem__(key)

    first_asset = object()
    second_asset = object()
    outlet_events = FakeOutletEvents()

    attach_output_row_count_metadata(
        outlet_events=outlet_events,
        airflow_assets=(first_asset, second_asset),
        metadata={
            "outputs": [
                {"asset_name": "finance.monthly_revrec_collections", "row_count": 2},
                {"asset_name": "finance.monthly_revrec_orders", "row_count": 3},
            ]
        },
    )

    assert outlet_events[first_asset].extra == {"row_count": 2}
    assert outlet_events[second_asset].extra == {"row_count": 3}


def test_python_asset_runner_metadata_follows_contract_output_order(
    tmp_path: Path,
    capsys,
) -> None:
    fake_runner = tmp_path / "fake-python-runner"
    fake_runner.write_text(
        "#!/usr/bin/env python3\n"
        "print('ordinary job log line')\n"
        "print('{\"row_counts\": [2, 3]}')\n"
    )
    fake_runner.chmod(0o755)

    config = replace(
        load_python_asset_job_config(PYTHON_REVREC_CONFIG_PATH, repo_root=REPO_ROOT),
        runner_executable_path=fake_runner,
    )

    assert run_python_asset_job(config) == {
        "outputs": [
            {
                "asset_name": "finance.monthly_revrec_collections",
                "asset_uri": "x-datalake://finance/monthly_revrec_collections",
                "row_count": 2,
            },
            {
                "asset_name": "finance.monthly_revrec_orders",
                "asset_uri": "x-datalake://finance/monthly_revrec_orders",
                "row_count": 3,
            },
        ]
    }
    assert "ordinary job log line" in capsys.readouterr().out


def test_python_asset_runner_skip_signal_skips_airflow_task(
    tmp_path: Path,
    capsys,
) -> None:
    from airflow.sdk.exceptions import AirflowSkipException

    fake_runner = tmp_path / "fake-python-runner"
    fake_runner.write_text(
        "#!/usr/bin/env python3\n"
        "print('ordinary job log line')\n"
        "print('{\"skip_reason\": \"no output products\"}')\n"
    )
    fake_runner.chmod(0o755)

    config = replace(
        load_python_asset_job_config(PYTHON_REVREC_CONFIG_PATH, repo_root=REPO_ROOT),
        runner_executable_path=fake_runner,
    )

    with pytest.raises(AirflowSkipException, match="no output products"):
        run_python_asset_job(config)
    assert "ordinary job log line" in capsys.readouterr().out


def test_python_asset_subprocess_serializes_no_output_products_skip(tmp_path: Path) -> None:
    fake_job = tmp_path / "fake_job.py"
    fake_job.write_text(
        "class NoOutputProductsError(RuntimeError):\n"
        "    pass\n\n"
        "def run():\n"
        "    raise NoOutputProductsError('no rows after filter')\n"
    )
    payload = json.dumps({"callable_path": "fake_job:run", "run_kwargs": {}})
    env = {**os.environ, "PYTHONPATH": f"{tmp_path}{os.pathsep}{REPO_ROOT}"}

    result = subprocess.run(
        [
            sys.executable,
            str(REPO_ROOT / "orchestration" / "airflow" / "engine" / "run_python_asset_job.py"),
            "--payload-json",
            payload,
        ],
        check=True,
        capture_output=True,
        env=env,
        text=True,
    )

    assert json.loads(result.stdout.splitlines()[-1]) == {"skip_reason": "no rows after filter"}


def test_python_revrec_dag_is_thin_contract_asset_registration() -> None:
    dag_source = (REPO_ROOT / "orchestration" / "airflow" / "dags" / "python_revrec_assets.py").read_text()

    assert "load_python_asset_job_config(PYTHON_REVREC_CONFIG_PATH)" in dag_source
    assert "schedule=PYTHON_REVREC_INLETS" in dag_source
    assert "inlets=PYTHON_REVREC_INLETS" in dag_source
    assert "outlets=PYTHON_REVREC_OUTLETS" in dag_source
    assert "execute_python_asset_job(PYTHON_REVREC)" in dag_source
    assert "attach_output_row_count_metadata(" in dag_source
    assert "x-datalake://finance/monthly_revrec_orders" not in dag_source


def test_monthly_revrec_process_empty_outputs_have_contract_columns() -> None:
    pd = pytest.importorskip("pandas")

    import transform.python.jobs.monthly_cloud_billing_revrec as job

    empty_orders = pd.DataFrame(columns=["order_status", "order_product_family"])
    empty_billing = pd.DataFrame(
        columns=[
            "total_amount",
            "segment_a_start_date",
            "segment_a_end_date",
            "segment_a_amount",
            "segment_b_start_date",
            "segment_b_end_date",
            "segment_b_amount",
            "segment_c_start_date",
            "segment_c_end_date",
            "segment_c_amount",
        ]
    )

    collections, orders = job.process(df_orders=empty_orders, df_billing=empty_billing)

    assert list(collections.columns) == job.COLLECTION_OUTPUT_COLUMNS
    assert list(orders.columns) == job.ORDER_OUTPUT_COLUMNS


def test_monthly_revrec_run_raises_no_output_when_no_billing_rows(monkeypatch: pytest.MonkeyPatch) -> None:
    pd = pytest.importorskip("pandas")

    import transform.python.jobs.monthly_cloud_billing_revrec as job

    def fake_read_df(_cur, sql: str, **_kwargs):
        if "salesforce_org_monthly_revrec" in sql:
            return pd.DataFrame(columns=["actual_revenue_recognition_date"])
        return pd.DataFrame(columns=["date_day", "order_service_start_date"])

    writes = []
    monkeypatch.setattr(job, "connect", lambda: object())
    monkeypatch.setattr(job, "read_df", fake_read_df)
    monkeypatch.setattr(job, "replace_table", lambda *args, **kwargs: writes.append((args, kwargs)))

    with pytest.raises(job.NoOutputProductsError, match="no billing rows"):
        job.run(billing_month="2026-01", order_date=date(2026, 6, 1))
    assert writes == []


def test_python_asset_job_shim_uses_pyjobs_extra_and_optional_env_file(tmp_path: Path) -> None:
    shim = REPO_ROOT / "orchestration" / "airflow" / "bin" / "python-asset-job-shim"
    fake_bin = tmp_path / "bin"
    fake_bin.mkdir()
    fake_uv = fake_bin / "uv"
    fake_uv.write_text(
        "#!/usr/bin/env bash\n"
        "printf 'UV_PROJECT_ENVIRONMENT=%s\\n' \"$UV_PROJECT_ENVIRONMENT\"\n"
        "printf 'PYTHONUNBUFFERED=%s\\n' \"$PYTHONUNBUFFERED\"\n"
        "printf '%s\\n' \"$@\"\n"
    )
    fake_uv.chmod(0o755)

    runtime_root = tmp_path / "runtime-root"
    runtime_root.mkdir()
    env = {
        **os.environ,
        "DATALAKE_AIRFLOW_REPO_ROOT": str(runtime_root),
        "PATH": f"{fake_bin}{os.pathsep}{os.environ['PATH']}",
    }

    without_env = subprocess.run(
        [str(shim), "--payload-json", "{}"],
        check=True,
        capture_output=True,
        env=env,
        text=True,
    ).stdout.splitlines()
    assert "UV_PROJECT_ENVIRONMENT=/opt/airflow/.venvs/pingcap-datalake-pyjobs" in without_env
    assert "PYTHONUNBUFFERED=1" in without_env
    assert "--extra" in without_env
    assert "pyjobs" in without_env
    assert "--env-file" not in without_env

    (runtime_root / ".env").write_text("DATABEND_DSN=redacted\n")
    with_env = subprocess.run(
        [str(shim), "--payload-json", "{}"],
        check=True,
        capture_output=True,
        env=env,
        text=True,
    ).stdout.splitlines()
    assert "--env-file" in with_env
    assert str(runtime_root / ".env") in with_env


def test_python_asset_runner_has_no_airflow_engine_import() -> None:
    runner_source = (
        REPO_ROOT / "orchestration" / "airflow" / "engine" / "run_python_asset_job.py"
    ).read_text()

    assert "orchestration.airflow.engine" not in runner_source


def test_python_revrec_dag_renders_contract_assets() -> None:
    import pytest

    airflow_models = pytest.importorskip("airflow.models")

    bag = airflow_models.DagBag(
        dag_folder=str(REPO_ROOT / "orchestration" / "airflow" / "dags" / "python_revrec_assets.py"),
        include_examples=False,
        safe_mode=False,
    )
    assert bag.import_errors == {}

    dag = bag.dags["orchestration_python_revrec_assets"]
    task = dag.get_task("monthly_cloud_billing_revrec")

    assert {asset.name for asset in dag.schedule} == {
        "finance.salesforce_order_daily_view",
        "finance.salesforce_org_monthly_revrec",
    }
    assert {asset.name for asset in task.inlets} == {
        "finance.salesforce_order_daily_view",
        "finance.salesforce_org_monthly_revrec",
    }
    assert {asset.name for asset in task.outlets} == {
        "finance.monthly_revrec_collections",
        "finance.monthly_revrec_orders",
    }


def test_build_python_airflow_asset_preserves_contract_group() -> None:
    import pytest

    pytest.importorskip("airflow.sdk")

    asset = build_airflow_asset(
        AirflowAssetSpec(
            name="finance.monthly_revrec_orders",
            uri="x-datalake://finance/monthly_revrec_orders",
            group="curated",
            extra={"contract_id": "finance.monthly_revrec_orders"},
        )
    )

    assert asset.group == "curated"