"""Tests for BatchRunner — mock-based, no network calls."""

from __future__ import annotations

from unittest.mock import MagicMock

from dalysdk.batch import BatchRunner

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _mock_client() -> MagicMock:
    """Return a mock DalyClient with energy_models and tasks stubs."""
    client = MagicMock()
    client.energy_models.create.return_value = {
        "taskDetails": {"taskId": 1, "status": "queued"},
        "energyModelId": 10,
        "runIndex": 1,
    }
    client.tasks.get.return_value = {"status": "RUNNING"}
    return client


# ---------------------------------------------------------------------------
# Submission
# ---------------------------------------------------------------------------


class TestSubmit:
    def test_submit_tracks_run(self):
        client = _mock_client()
        batch = BatchRunner(client)

        resp = batch.submit("Site A", {"blocks": []})

        assert resp["taskDetails"]["taskId"] == 1
        assert len(batch.runs) == 1
        run = batch.runs[0]
        assert run["name"] == "Site A"
        assert run["task_id"] == 1
        assert run["energy_model_id"] == 10
        assert run["status"] == "PENDING"

    def test_submit_passes_timeout(self):
        client = _mock_client()
        batch = BatchRunner(client)

        batch.submit("Site A", {"blocks": []}, timeout=120.0)

        client.energy_models.create.assert_called_once_with({"blocks": []}, timeout=120.0)

    def test_submit_many_sequential(self):
        client = _mock_client()
        # Give each call a different task_id
        client.energy_models.create.side_effect = [
            {"taskDetails": {"taskId": 1}, "energyModelId": 10, "runIndex": 1},
            {"taskDetails": {"taskId": 2}, "energyModelId": 11, "runIndex": 2},
            {"taskDetails": {"taskId": 3}, "energyModelId": 12, "runIndex": 3},
        ]
        batch = BatchRunner(client)

        count = batch.submit_many(
            [
                ("A", {"blocks": []}),
                ("B", {"blocks": []}),
                ("C", {"blocks": []}),
            ]
        )

        assert count == 3
        assert len(batch.runs) == 3
        assert [r["name"] for r in batch.runs] == ["A", "B", "C"]

    def test_submit_many_parallel(self):
        client = _mock_client()
        client.energy_models.create.side_effect = [
            {"taskDetails": {"taskId": i}, "energyModelId": 10 + i, "runIndex": i} for i in range(1, 4)
        ]
        batch = BatchRunner(client, max_workers=2)

        count = batch.submit_many(
            [("A", {}), ("B", {}), ("C", {})],
            parallel=True,
        )

        assert count == 3
        assert len(batch.runs) == 3

    def test_submit_failure_counted(self):
        client = _mock_client()
        client.energy_models.create.side_effect = [
            {"taskDetails": {"taskId": 1}, "energyModelId": 10, "runIndex": 1},
            RuntimeError("API down"),
            {"taskDetails": {"taskId": 3}, "energyModelId": 12, "runIndex": 3},
        ]
        batch = BatchRunner(client)

        count = batch.submit_many([("A", {}), ("B", {}), ("C", {})])

        assert count == 2
        assert len(batch.runs) == 2

    def test_submit_accepts_legacy_id_fields(self):
        client = _mock_client()
        client.energy_models.create.return_value = {
            "task_id": 41,
            "energy_model_id": 77,
            "run_index": 8,
        }
        batch = BatchRunner(client)

        batch.submit("Legacy", {"blocks": []})
        run = batch.runs[0]
        assert run["task_id"] == 41
        assert run["energy_model_id"] == 77
        assert run["run_index"] == 8


# ---------------------------------------------------------------------------
# Polling
# ---------------------------------------------------------------------------


class TestPoll:
    def test_poll_updates_status(self):
        client = _mock_client()
        batch = BatchRunner(client)
        batch.submit("Site A", {})

        # First poll: still running
        client.tasks.get.return_value = {"status": "RUNNING"}
        c, f, p = batch.poll()
        assert (c, f, p) == (0, 0, 1)

        # Second poll: completed
        client.tasks.get.return_value = {"status": "COMPLETED"}
        c, f, p = batch.poll()
        assert (c, f, p) == (1, 0, 0)

    def test_poll_marks_failed(self):
        client = _mock_client()
        batch = BatchRunner(client)
        batch.submit("Site A", {})

        client.tasks.get.return_value = {"status": "FAILED", "error": "OOM"}
        c, f, p = batch.poll()
        assert (c, f, p) == (0, 1, 0)
        assert batch.failed[0]["error"] == "OOM"

    def test_poll_marks_failed_with_task_details_error(self):
        client = _mock_client()
        batch = BatchRunner(client)
        batch.submit("Site A", {})

        client.tasks.get.return_value = {"status": "FAILED", "taskDetails": {"errorMessage": "Validation failed"}}
        c, f, p = batch.poll()
        assert (c, f, p) == (0, 1, 0)
        assert batch.failed[0]["error"] == "Validation failed"

    def test_poll_skips_terminal(self):
        """Terminal runs should not be re-polled."""
        client = _mock_client()
        batch = BatchRunner(client)
        batch.submit("Site A", {})

        # Complete it
        client.tasks.get.return_value = {"status": "COMPLETED"}
        batch.poll()
        client.tasks.get.reset_mock()

        # Next poll should not call tasks.get again
        batch.poll()
        client.tasks.get.assert_not_called()


# ---------------------------------------------------------------------------
# Wait
# ---------------------------------------------------------------------------


class TestWait:
    def test_wait_completes(self):
        client = _mock_client()
        batch = BatchRunner(client)
        batch.submit("Site A", {})

        # tasks.get returns COMPLETED on first poll
        client.tasks.get.return_value = {"status": "COMPLETED"}

        result = batch.wait(poll_interval=0.01)
        assert result is True
        assert len(batch.completed) == 1

    def test_wait_timeout(self):
        client = _mock_client()
        batch = BatchRunner(client)
        batch.submit("Site A", {})

        # Always return RUNNING so it never finishes
        client.tasks.get.return_value = {"status": "RUNNING"}

        result = batch.wait(timeout=0.05, poll_interval=0.01)
        assert result is False
        assert len(batch.pending) == 1

    def test_on_poll_callback(self):
        client = _mock_client()
        batch = BatchRunner(client)
        batch.submit("Site A", {})

        client.tasks.get.return_value = {"status": "COMPLETED"}

        calls: list[tuple[int, int, int]] = []
        batch.wait(poll_interval=0.01, on_poll=lambda c, f, p: calls.append((c, f, p)))

        assert len(calls) >= 1
        assert calls[-1] == (1, 0, 0)


# ---------------------------------------------------------------------------
# Summary / Results
# ---------------------------------------------------------------------------


class TestSummary:
    def test_summary_structure(self):
        client = _mock_client()
        client.energy_models.create.side_effect = [
            {"taskDetails": {"taskId": 1}, "energyModelId": 10, "runIndex": 1},
            {"taskDetails": {"taskId": 2}, "energyModelId": 11, "runIndex": 2},
        ]
        batch = BatchRunner(client)
        batch.submit("A", {})
        batch.submit("B", {})

        # Complete A, fail B
        def get_task(tid: int) -> dict:
            if tid == 1:
                return {"status": "COMPLETED"}
            return {"status": "FAILED", "error": "bad input"}

        client.tasks.get.side_effect = get_task
        batch.poll()

        s = batch.summary()
        assert s["total"] == 2
        assert s["completed"] == 1
        assert s["failed"] == 1
        assert s["pending"] == 0
        assert len(s["results"]) == 1
        assert s["results"][0]["name"] == "A"
        assert len(s["failures"]) == 1
        assert s["failures"][0]["name"] == "B"
        assert s["failures"][0]["error"] == "bad input"


class TestResults:
    def test_results_fetches_models(self):
        client = _mock_client()
        batch = BatchRunner(client)
        batch.submit("Site A", {})

        client.tasks.get.return_value = {"status": "COMPLETED"}
        batch.poll()

        client.energy_models.get.return_value = {
            "id": 10,
            "projectId": 4,
            "globalProjectId": 4004,
            "annualProductionMWh": 123.4,
            "revision": "v0001",
        }

        results = batch.results()
        assert len(results) == 1
        assert results[0]["name"] == "Site A"
        assert results[0]["data"]["annualProductionMWh"] == 123.4
        assert results[0]["data"]["globalProjectId"] == 4004
        client.energy_models.get.assert_called_once_with(10)

    def test_results_skips_none_id(self):
        client = _mock_client()
        client.energy_models.create.return_value = {
            "taskDetails": {"taskId": 1},
            "energyModelId": None,
            "runIndex": None,
        }
        batch = BatchRunner(client)
        batch.submit("Site A", {})

        client.tasks.get.return_value = {"status": "COMPLETED"}
        batch.poll()

        results = batch.results()
        assert len(results) == 0
        client.energy_models.get.assert_not_called()


# ---------------------------------------------------------------------------
# Properties
# ---------------------------------------------------------------------------


class TestProperties:
    def test_filtered_views(self):
        client = _mock_client()
        client.energy_models.create.side_effect = [
            {"taskDetails": {"taskId": i}, "energyModelId": 10 + i, "runIndex": i} for i in range(1, 4)
        ]
        batch = BatchRunner(client)
        batch.submit("A", {})
        batch.submit("B", {})
        batch.submit("C", {})

        def get_task(tid: int) -> dict:
            return {
                1: {"status": "COMPLETED"},
                2: {"status": "FAILED"},
                3: {"status": "RUNNING"},
            }[tid]

        client.tasks.get.side_effect = get_task
        batch.poll()

        assert len(batch.completed) == 1
        assert len(batch.failed) == 1
        assert len(batch.pending) == 1
        assert batch.completed[0]["name"] == "A"
        assert batch.failed[0]["name"] == "B"
        assert batch.pending[0]["name"] == "C"
