import numpy as np
import xarray as xr

from dr_gcm.download_ecmwf_asia import load_domain
from dr_gcm.download_ecmwf_pilot import (
    load_pilot_config,
    request_for_group,
    retrieve_group,
    run_pilot,
    validate_pilot_file,
)


def test_pilot_requests_are_small_and_split_by_semantics():
    domain = load_domain()
    pilot = load_pilot_config()
    requests = {name: request_for_group("1981-01", domain, pilot, name) for name in pilot["groups"]}
    assert set(requests) == {"instantaneous", "accumulated", "soil_moisture"}
    assert requests["instantaneous"]["leadtime_hour"] == ["0", "6", "24", "30"]
    assert requests["instantaneous"]["area"] == [40, 70, -20, 150]
    assert requests["accumulated"]["leadtime_hour"] == ["24", "48", "72"]
    assert requests["accumulated"]["variable"] == [
        "total_precipitation", "surface_latent_heat_flux", "surface_sensible_heat_flux",
        "top_net_solar_radiation", "top_net_thermal_radiation",
    ]
    assert requests["soil_moisture"]["variable"] == ["volumetric_soil_moisture"]
    assert requests["soil_moisture"]["leadtime_hour"] == ["24", "48", "72"]


def test_pilot_validation_requires_variables_and_25_members(tmp_path):
    path = tmp_path / "pilot.nc"
    xr.Dataset(
        {"tp": (("number", "forecast_period", "latitude", "longitude"), np.zeros((25, 2, 1, 1)))},
        coords={"number": np.arange(25), "forecast_period": [6, 24], "latitude": [10.0], "longitude": [105.0]},
    ).to_netcdf(path)
    report = validate_pilot_file(path, ["tp"])
    assert report["status"] == "ok"
    validate_pilot_file(path, ["tp"], expected_leadtime_hour=["6", "24"])
    try:
        validate_pilot_file(path, ["tp", "slhf"])
    except ValueError as exc:
        assert "slhf" in str(exc)
    else:
        raise AssertionError("missing expected variable was not rejected")
    try:
        validate_pilot_file(path, ["tp"], expected_leadtime_hour=["24", "48"])
    except ValueError as exc:
        assert "Expected lead hours" in str(exc)
    else:
        raise AssertionError("unexpected lead hours were not rejected")


def test_pilot_validation_requires_four_soil_layers(tmp_path):
    path = tmp_path / "soil.nc"
    xr.Dataset(
        {
            "vsw": (
                ("number", "forecast_period", "soilLayer", "latitude", "longitude"),
                np.zeros((25, 3, 4, 1, 1)),
            )
        },
        coords={
            "number": np.arange(25),
            "forecast_period": [24, 48, 72],
            "soilLayer": [1, 2, 3, 4],
            "latitude": [10.0],
            "longitude": [105.0],
        },
    ).to_netcdf(path)
    validate_pilot_file(
        path, ["vsw"], expected_leadtime_hour=["24", "48", "72"],
        expected_soil_layers=4,
    )


def test_pilot_retry_partial_replaces_only_incomplete_file(tmp_path):
    domain = load_domain()
    pilot = load_pilot_config()
    partial = tmp_path / "ecmwf_seas51_pilot_1981-01_soil_moisture.nc.part"
    partial.write_text("incomplete", encoding="utf-8")

    class Client:
        def retrieve(self, dataset, request, target):
            xr.Dataset(
                {
                    "vsw": (
                        ("number", "forecast_period", "soilLayer", "latitude", "longitude"),
                        np.zeros((25, 3, 4, 1, 1)),
                    )
                },
                coords={
                    "number": np.arange(25),
                    "forecast_period": [24, 48, 72],
                    "soilLayer": [1, 2, 3, 4],
                    "latitude": [10.0],
                    "longitude": [105.0],
                },
            ).to_netcdf(target)

    result = retrieve_group(
        Client(), "1981-01", "soil_moisture", domain, pilot, tmp_path,
        retry_partial=True,
    )
    assert result["status"] == "downloaded"
    assert not partial.exists()


def test_pilot_checkpoints_manifest_after_each_group(tmp_path):
    domain = load_domain()
    pilot = load_pilot_config()

    class InterruptedClient:
        calls = 0

        def retrieve(self, dataset, request, target):
            self.calls += 1
            if self.calls == 1:
                xr.Dataset(
                    {
                        name: (
                            ("number", "forecast_period", "latitude", "longitude"),
                            np.zeros((25, 3, 1, 1)),
                        )
                        for name in pilot["groups"]["instantaneous"]["expected_netcdf_variables"]
                    },
                    coords={
                        "number": np.arange(25),
                        "forecast_period": [6, 24, 30],
                        "latitude": [10.0],
                        "longitude": [105.0],
                    },
                ).to_netcdf(target)
            else:
                raise KeyboardInterrupt

    try:
        run_pilot(
            "1981-01",
            ["instantaneous", "accumulated"],
            pilot,
            domain,
            tmp_path,
            client=InterruptedClient(),
        )
    except KeyboardInterrupt:
        pass
    else:
        raise AssertionError("simulated interruption was not raised")

    manifest = __import__("json").loads(
        (tmp_path / "ecmwf_seas51_pilot_1981-01_manifest.json").read_text(encoding="utf-8")
    )
    assert [result["status"] for result in manifest["results"]] == ["downloaded"]