from pathlib import Path

import numpy as np
import pytest
import xarray as xr

from dr_gcm.recover_ecmwf_instantaneous import (
    ExternalDataUnavailableError,
    damaged_tape_error,
    merge_chunks,
    reset_blocked_attempt_cycle,
    retrieve_adaptive,
    split_chunks,
)


def domain():
    return {
        "originating_centre": "ecmwf",
        "system": "51",
        "area": [40, 70, -20, 150],
        "download_groups": {
            "instantaneous": {
                "expected_members": 25,
            }
        },
    }


def test_split_chunks_and_damaged_tape_detection():
    assert split_chunks(["6", "12", "18", "24", "30"], 2) == [
        ["6", "12"], ["18", "24"], ["30"]
    ]
    with pytest.raises(ValueError):
        split_chunks(["6"], 0)
    assert damaged_tape_error(RuntimeError("AccessError: damaged tape K0325300"))
    assert not damaged_tape_error(RuntimeError("temporary network failure"))


def test_adaptive_retrieval_splits_and_resumes(tmp_path):
    class Client:
        def __init__(self):
            self.calls = []

        def retrieve(self, dataset, request, target):
            leads = request["leadtime_hour"]
            self.calls.append(tuple(leads))
            if len(leads) > 1:
                raise RuntimeError("AccessError: requested data is on damaged tape")
            Path(target).write_text("chunk", encoding="utf-8")

    def validator(path, init, spec, leads):
        assert path.exists()
        return {"lead_count": len(leads)}

    manifest = {"chunks": {}}
    manifest_path = tmp_path / "manifest.json"
    client = Client()
    paths = retrieve_adaptive(
        client, "dataset", "1989-03", domain(), "stl1",
        "soil_temperature_level_1", ["6", "12", "18", "24"],
        tmp_path / "work", manifest, manifest_path, validator,
    )
    assert len(paths) == 4
    assert len(client.calls) == 6
    assert all(path.exists() for path in paths)

    resumed = Client()
    second = retrieve_adaptive(
        resumed, "dataset", "1989-03", domain(), "stl1",
        "soil_temperature_level_1", ["6", "12", "18", "24"],
        tmp_path / "work", manifest, manifest_path, validator,
    )
    assert second == paths
    assert resumed.calls == []


def test_adaptive_retrieval_does_not_split_unrelated_errors(tmp_path):
    class Client:
        def retrieve(self, dataset, request, target):
            raise RuntimeError("authentication failed")

    with pytest.raises(RuntimeError, match="authentication"):
        retrieve_adaptive(
            Client(), "dataset", "1989-03", domain(), "stl1",
            "soil_temperature_level_1", ["6", "12"], tmp_path / "work",
            {"chunks": {}}, tmp_path / "manifest.json", lambda *args: {},
        )


def test_single_lead_damaged_tape_has_circuit_breaker(tmp_path):
    class Client:
        def __init__(self):
            self.calls = 0

        def retrieve(self, dataset, request, target):
            self.calls += 1
            raise RuntimeError("AccessError: requested data is on damaged tape K0325300")

    manifest = {"chunks": {}}
    manifest_path = tmp_path / "manifest.json"
    client = Client()
    for expected_attempt in (1, 2):
        with pytest.raises(RuntimeError, match="damaged tape"):
            retrieve_adaptive(
                client, "dataset", "1989-03", domain(), "stl1",
                "soil_temperature_level_1", ["750"], tmp_path / "work",
                manifest, manifest_path, lambda *args: {},
                max_single_step_attempts=3,
            )
        assert manifest["chunks"]["stl1:750-750:1"]["attempt_count"] == expected_attempt

    with pytest.raises(ExternalDataUnavailableError, match="after 3 attempts"):
        retrieve_adaptive(
            client, "dataset", "1989-03", domain(), "stl1",
            "soil_temperature_level_1", ["750"], tmp_path / "work",
            manifest, manifest_path, lambda *args: {},
            max_single_step_attempts=3,
        )
    assert client.calls == 3

    with pytest.raises(ExternalDataUnavailableError, match="after 3 attempts"):
        retrieve_adaptive(
            client, "dataset", "1989-03", domain(), "stl1",
            "soil_temperature_level_1", ["750"], tmp_path / "work",
            manifest, manifest_path, lambda *args: {},
            max_single_step_attempts=3,
        )
    assert client.calls == 3


def test_new_attempt_cycle_retries_blocked_single_lead_with_provenance(tmp_path):
    class Client:
        def __init__(self):
            self.calls = 0

        def retrieve(self, dataset, request, target):
            self.calls += 1
            Path(target).write_text("chunk", encoding="utf-8")

    key = "stl1:750-750:1"
    previous = {
        "status": "blocked_external",
        "updated_utc": "2026-09-30T00:00:00+00:00",
        "blocked_reason": "damaged tape after 3 attempts",
        "chunks": {
            key: {
                "status": "error",
                "leads": ["750"],
                "attempt_count": 3,
                "error": "AccessError: requested data is on damaged tape",
            }
        },
    }
    manifest = {"chunks": previous["chunks"], "single_step_variables": ["stl1"]}
    assert reset_blocked_attempt_cycle(manifest, previous) == [key]
    assert manifest["chunks"][key]["attempt_count"] == 0
    assert manifest["chunks"][key]["total_attempt_count"] == 3

    client = Client()
    paths = retrieve_adaptive(
        client, "dataset", "1989-03", domain(), "stl1",
        "soil_temperature_level_1", ["750"], tmp_path / "work",
        manifest, tmp_path / "manifest.json", lambda *args: {"lead_count": 1},
        max_single_step_attempts=3,
    )
    assert client.calls == 1
    assert len(paths) == 1
    assert manifest["chunks"][key]["attempt_count"] == 1
    assert manifest["chunks"][key]["total_attempt_count"] == 4
    assert len(manifest["recovery_cycles"]) == 1


def _write_chunk(path, variable, leads, offset):
    periods = np.asarray(leads, dtype="timedelta64[h]")
    values = np.arange(len(leads) * 2, dtype=np.float32).reshape(len(leads), 2, 1, 1)
    values += offset
    reference = np.datetime64("1989-03-01T00:00:00", "ns")
    xr.Dataset(
        {
            variable: (
                ("forecast_period", "number", "latitude", "longitude"),
                values,
            )
        },
        coords={
            "forecast_period": periods,
            "number": [0, 1],
            "latitude": [20.0],
            "longitude": [105.0],
            "forecast_reference_time": reference,
            "valid_time": ("forecast_period", reference + periods),
        },
    ).to_netcdf(path)


def test_merge_chunks_requires_complete_leads_and_promotes_atomically(tmp_path):
    chunk_paths = {}
    for variable_index, variable in enumerate(("t2m", "sst")):
        first = tmp_path / f"{variable}_first.nc"
        second = tmp_path / f"{variable}_second.nc"
        _write_chunk(first, variable, [6, 12], variable_index * 100)
        _write_chunk(second, variable, [18, 24], variable_index * 100)
        chunk_paths[variable] = [second, first]
    target = tmp_path / "instantaneous.nc"
    spec = {
        "expected_netcdf_variables": ["t2m", "sst"],
        "expected_members": 2,
        "allow_missing_initial_step": True,
    }
    report = merge_chunks(
        chunk_paths, ["6", "12", "18", "24"], target, "1989-03", spec
    )
    assert report["lead_count"] == 4
    assert target.exists()
    assert not target.with_suffix(".nc.part").exists()
    with xr.open_dataset(target) as dataset:
        assert list(dataset.data_vars) == ["t2m", "sst"]
        assert dataset.sizes["forecast_period"] == 4

    missing = dict(chunk_paths)
    missing["t2m"] = [chunk_paths["t2m"][0]]
    with pytest.raises(ValueError, match="lead mismatch"):
        merge_chunks(
            missing, ["6", "12", "18", "24"], tmp_path / "bad.nc",
            "1989-03", spec,
        )