from __future__ import annotations

import unittest
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


PACKAGE_DIR = Path(__file__).resolve().parents[1] / "modelhub_submmit_api"
if str(PACKAGE_DIR) in sys.path:
    sys.path.remove(str(PACKAGE_DIR))
sys.path.insert(0, str(PACKAGE_DIR))

from architecture_compatibility import architecture_compatibility_key  # noqa: E402
from modelhub_client import ModelHubClient, ModelHubClientPool  # noqa: E402
from queue_cleanup import (  # noqa: E402
    OwnedTask,
    cleanup_certain_oom_tasks,
    find_architecture_incompatible_tasks,
    find_certain_oom_tasks,
    find_old_overflow_tasks,
)


GIB = 1024**3


class FakeQueueClient:
    def __init__(
        self,
        records: list[dict[str, Any]],
        *,
        disappear_on_recheck: bool = False,
        drop_first_on_recheck: bool = False,
        promote_on_waiting_read: int | None = None,
    ) -> None:
        self.records = list(records)
        self.disappear_on_recheck = disappear_on_recheck
        self.drop_first_on_recheck = drop_first_on_recheck
        self.promote_on_waiting_read = promote_on_waiting_read
        self.waiting_reads = 0
        self.stopped: list[list[int]] = []

    def list_tasks_page(self, *, status: str, **_kwargs: Any) -> dict[str, Any]:
        if status == "waiting":
            self.waiting_reads += 1
            if self.promote_on_waiting_read == self.waiting_reads:
                waiting = [record for record in self.records if record["status"] == "waiting"]
                if waiting:
                    max(waiting, key=lambda item: int(item["taskId"]))["status"] = "running"
            if self.disappear_on_recheck and self.waiting_reads >= 2:
                records: list[dict[str, Any]] = []
            else:
                records = [record for record in self.records if record["status"] == "waiting"]
                if self.drop_first_on_recheck and self.waiting_reads >= 2:
                    records = sorted(records, key=lambda item: int(item["taskId"]))[1:]
        else:
            records = [record for record in self.records if record["status"] == status]
        return {"code": 0, "data": {"records": records, "pages": 1}}

    def stop_tasks(self, task_ids: list[int]) -> dict[str, Any]:
        self.stopped.append(list(task_ids))
        ids = set(task_ids)
        self.records = [record for record in self.records if int(record["taskId"]) not in ids]
        return {"code": 0, "data": None}

    def count_active_tasks(self, **_kwargs: Any) -> int:
        return len(self.records)


class FakeDiscovery:
    def __init__(
        self,
        sizes: dict[str, int | None],
        last_modified: dict[str, datetime | None] | None = None,
        configs: dict[str, dict[str, Any]] | None = None,
    ) -> None:
        self.sizes = sizes
        self.last_modified = last_modified or {}
        self.configs = configs or {}

    def list_repo_tree(self, repo_id: str) -> list[dict[str, Any]]:
        size = self.sizes[repo_id]
        if size is None:
            return [{"Path": "model.safetensors"}]
        return [{"Path": "model.safetensors", "Size": size}]

    def get_model_last_modified(self, repo_id: str) -> datetime | None:
        return self.last_modified.get(repo_id)

    def get_model_config(self, repo_id: str) -> tuple[dict[str, Any], str | None]:
        config = self.configs.get(repo_id)
        return (dict(config), None) if config is not None else ({}, "config_not_found")


class RecordingHttpClient:
    def __init__(self) -> None:
        self.calls: list[tuple[str, str, dict[str, Any] | None, dict[str, Any] | None]] = []

    def request_json(
        self,
        method: str,
        path: str,
        *,
        query: dict[str, Any] | None = None,
        data: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        self.calls.append((method, path, query, data))
        return {"code": 0, "data": None}


class QueueCleanupTests(unittest.TestCase):
    @staticmethod
    def architecture_block(
        *,
        gpu: str = "Iluvatar_bi-100",
        framework: str = "vllm",
        task_type: str = "text-generation",
        signature: str = "architectures:qwen2forcausallm",
    ) -> tuple[str, dict[str, Any]]:
        key = architecture_compatibility_key(gpu, framework, task_type, signature)
        assert key is not None
        return key, {
            "targetGpu": gpu,
            "framework": framework,
            "taskType": task_type,
            "matchType": "architectures",
            "architectureSignature": signature,
            "evidenceCount": 1,
            "expiresAt": "2026-09-11T00:00:00+00:00",
        }

    def test_architecture_cleanup_matches_exact_context_and_protects_running(self) -> None:
        key, block = self.architecture_block()
        tasks = [
            OwnedTask(0, 1, "owner/waiting", "Iluvatar_bi-100", "waiting"),
            OwnedTask(0, 2, "owner/running", "Iluvatar_bi-100", "running"),
            OwnedTask(0, 3, "owner/other-framework", "Iluvatar_bi-100", "waiting"),
        ]
        contexts = {
            "1": {
                "modelId": "owner/waiting",
                "targetGpu": "Iluvatar_bi-100",
                "framework": "vllm",
                "taskType": "text-generation",
                "modelProfile": {"architectures": ["Qwen2ForCausalLM"]},
            },
            "2": {
                "modelId": "owner/running",
                "targetGpu": "Iluvatar_bi-100",
                "framework": "vllm",
                "taskType": "text-generation",
                "modelProfile": {"architectures": ["Qwen2ForCausalLM"]},
            },
            "3": {
                "modelId": "owner/other-framework",
                "targetGpu": "Iluvatar_bi-100",
                "framework": "mindie",
                "taskType": "text-generation",
                "modelProfile": {"architectures": ["Qwen2ForCausalLM"]},
            },
        }

        selected, skipped = find_architecture_incompatible_tasks(
            tasks,
            architecture_blocks={key: block},
            task_contexts=contexts,
            model_configs={},
        )

        self.assertEqual([1], [item["taskId"] for item in selected])
        self.assertEqual(1, skipped["runningMatchedProtected"])
        self.assertEqual(1, skipped["noMatchingBlock"])

    def test_queue_cleanup_fetches_config_and_stops_known_incompatible_waiting_task(self) -> None:
        key, block = self.architecture_block()
        client = FakeQueueClient(
            [
                {
                    "taskId": 1,
                    "modelId": "owner/model",
                    "gpuType": "Iluvatar_bi-100",
                    "status": "waiting",
                }
            ]
        )
        pool = ModelHubClientPool([client], active_task_cap=100)  # type: ignore[list-item]
        summary = cleanup_certain_oom_tasks(
            pool,
            FakeDiscovery(
                {"owner/model": 1 * GIB},
                configs={
                    "owner/model": {
                        "model_type": "qwen2",
                        "architectures": ["Qwen2ForCausalLM"],
                    }
                },
            ),  # type: ignore[arg-type]
            architecture_compatibility_blocks={key: block},
            task_compatibility_contexts={
                "1": {
                    "modelId": "owner/model",
                    "targetGpu": "Iluvatar_bi-100",
                    "framework": "vllm",
                    "taskType": "text-generation",
                    "modelProfile": {},
                }
            },
            log=lambda _message: None,
        )

        self.assertEqual(1, summary["architectureIncompatibleCount"])
        self.assertEqual(1, summary["cancelledCount"])
        self.assertEqual([[1]], client.stopped)

    def test_dynamic_architecture_only_cleanup_skips_expensive_size_and_age_scans(self) -> None:
        key, block = self.architecture_block()
        client = FakeQueueClient(
            [
                {
                    "taskId": 1,
                    "modelId": "owner/model",
                    "gpuType": "Iluvatar_bi-100",
                    "status": "waiting",
                }
            ]
        )
        pool = ModelHubClientPool([client], active_task_cap=100)  # type: ignore[list-item]
        summary = cleanup_certain_oom_tasks(
            pool,
            FakeDiscovery({}),  # type: ignore[arg-type]
            architecture_compatibility_blocks={key: block},
            task_compatibility_contexts={
                "1": {
                    "modelId": "owner/model",
                    "targetGpu": "Iluvatar_bi-100",
                    "framework": "vllm",
                    "taskType": "text-generation",
                    "modelProfile": {"architectures": ["Qwen2ForCausalLM"]},
                }
            },
            architecture_only=True,
            log=lambda _message: None,
        )

        self.assertTrue(summary["architectureOnly"])
        self.assertEqual(0, summary["repositorySizesComplete"])
        self.assertEqual(0, summary["modelAgeMetadataComplete"])
        self.assertEqual(1, summary["architectureIncompatibleCount"])
        self.assertEqual([[1]], client.stopped)

    def test_architecture_cleanup_recheck_protects_task_that_started_running(self) -> None:
        key, block = self.architecture_block()
        client = FakeQueueClient(
            [
                {
                    "taskId": 1,
                    "modelId": "owner/model",
                    "gpuType": "Iluvatar_bi-100",
                    "status": "waiting",
                }
            ],
            promote_on_waiting_read=2,
        )
        pool = ModelHubClientPool([client], active_task_cap=100)  # type: ignore[list-item]
        summary = cleanup_certain_oom_tasks(
            pool,
            FakeDiscovery({"owner/model": 1 * GIB}),  # type: ignore[arg-type]
            architecture_compatibility_blocks={key: block},
            task_compatibility_contexts={
                "1": {
                    "modelId": "owner/model",
                    "targetGpu": "Iluvatar_bi-100",
                    "framework": "vllm",
                    "taskType": "text-generation",
                    "modelProfile": {"architectures": ["Qwen2ForCausalLM"]},
                }
            },
            read_concurrency=1,
            log=lambda _message: None,
        )

        self.assertEqual(1, summary["architectureIncompatibleCount"])
        self.assertEqual(0, summary["cancelledCount"])
        self.assertEqual("task_started_running", summary["policyNoLongerAppliesTasks"][0]["policyChangeReason"])
        self.assertEqual([], client.stopped)

    def test_old_models_use_each_accounts_own_capacity_minus_ten_threshold(self) -> None:
        now = datetime(2026, 8, 11, tzinfo=timezone.utc)
        tasks = [
            OwnedTask(0, index, "owner/old", "Iluvatar_bi-100", "waiting")
            for index in range(1, 92)
        ]
        tasks.extend(
            OwnedTask(1, 1000 + index, "owner/old", "Iluvatar_bi-100", "waiting")
            for index in range(1, 192)
        )
        tasks.append(OwnedTask(0, 92, "owner/recent", "Iluvatar_bi-100", "waiting"))
        tasks.append(OwnedTask(0, 93, "owner/old", "Iluvatar_bi-100", "running"))
        tasks.append(OwnedTask(1, 1192, "owner/recent", "Iluvatar_bi-100", "waiting"))

        selected, skipped = find_old_overflow_tasks(
            tasks,
            model_last_modified={
                "owner/old": datetime(2026, 7, 1, tzinfo=timezone.utc),
                "owner/recent": datetime(2026, 8, 10, tzinfo=timezone.utc),
            },
            queue_threshold={0: 90, 1: 190},
            recent_model_days=7,
            reference_time=now,
        )

        self.assertEqual([91, 1191], [item["taskId"] for item in selected])
        self.assertEqual([91, 191], [item["queuePosition"] for item in selected])
        self.assertEqual(2, skipped["recentOverflowTasks"])
        self.assertEqual(1, skipped["runningOverflowProtected"])

    def test_age_cleanup_never_stops_running_overflow_task(self) -> None:
        records = [
            {
                "taskId": index,
                "modelId": "owner/old",
                "gpuType": "Iluvatar_bi-100",
                "status": "running" if index == 91 else "waiting",
            }
            for index in range(1, 92)
        ]
        client = FakeQueueClient(records)
        pool = ModelHubClientPool(
            [client],  # type: ignore[list-item]
            active_task_cap=100,
            recent_model_reserve_slots=10,
        )
        summary = cleanup_certain_oom_tasks(
            pool,
            FakeDiscovery(
                {"owner/old": 1 * GIB},
                {"owner/old": datetime(2026, 7, 1, tzinfo=timezone.utc)},
            ),  # type: ignore[arg-type]
            reference_time=datetime(2026, 8, 11, tzinfo=timezone.utc),
            log=lambda _message: None,
        )

        self.assertEqual(0, summary["oldOverflowCount"])
        self.assertEqual(1, summary["agePolicySkipped"]["runningOverflowProtected"])
        self.assertEqual([], client.stopped)

    def test_age_cleanup_recheck_releases_task_that_started_running(self) -> None:
        records = [
            {
                "taskId": index,
                "modelId": "owner/old",
                "gpuType": "Iluvatar_bi-100",
                "status": "waiting",
            }
            for index in range(1, 92)
        ]
        # Read 1 is discovery, read 2 is the account-wide mutation recheck,
        # and read 3 is the final age-only recheck after the OOM phase.
        client = FakeQueueClient(records, promote_on_waiting_read=3)
        pool = ModelHubClientPool(
            [client],  # type: ignore[list-item]
            active_task_cap=100,
            recent_model_reserve_slots=10,
        )
        summary = cleanup_certain_oom_tasks(
            pool,
            FakeDiscovery(
                {"owner/old": 1 * GIB},
                {"owner/old": datetime(2026, 7, 1, tzinfo=timezone.utc)},
            ),  # type: ignore[arg-type]
            read_concurrency=1,
            reference_time=datetime(2026, 8, 11, tzinfo=timezone.utc),
            log=lambda _message: None,
        )

        self.assertEqual(1, summary["oldOverflowCount"])
        self.assertEqual(0, summary["cancelledCount"])
        self.assertEqual("task_started_running", summary["policyNoLongerAppliesTasks"][0]["policyChangeReason"])
        self.assertEqual([], client.stopped)

    def test_certain_oom_cleanup_can_still_stop_running_task(self) -> None:
        client = FakeQueueClient(
            [
                {
                    "taskId": 1,
                    "modelId": "owner/large",
                    "gpuType": "Iluvatar_bi-100",
                    "status": "running",
                }
            ]
        )
        pool = ModelHubClientPool([client], active_task_cap=100)  # type: ignore[list-item]
        summary = cleanup_certain_oom_tasks(
            pool,
            FakeDiscovery({"owner/large": 40 * GIB}),  # type: ignore[arg-type]
            log=lambda _message: None,
        )

        self.assertEqual(1, summary["certainOomCount"])
        self.assertEqual(1, summary["cancelledCount"])
        self.assertEqual([[1]], client.stopped)

    def test_old_overflow_task_is_not_stopped_if_it_moves_inside_dynamic_threshold(self) -> None:
        records = [
            {
                "taskId": index,
                "modelId": "owner/old",
                "gpuType": "Iluvatar_bi-100",
                "status": "waiting",
            }
            for index in range(1, 92)
        ]
        client = FakeQueueClient(records, drop_first_on_recheck=True)
        pool = ModelHubClientPool(
            [client],  # type: ignore[list-item]
            active_task_cap=100,
            recent_model_reserve_slots=10,
        )
        summary = cleanup_certain_oom_tasks(
            pool,
            FakeDiscovery(
                {"owner/old": 1 * GIB},
                {"owner/old": datetime(2026, 7, 1, tzinfo=timezone.utc)},
            ),  # type: ignore[arg-type]
            reference_time=datetime(2026, 8, 11, tzinfo=timezone.utc),
            log=lambda _message: None,
        )

        self.assertEqual(1, summary["oldOverflowCount"])
        self.assertEqual(0, summary["cancelledCount"])
        self.assertEqual(1, summary["policyNoLongerAppliesCount"])
        self.assertEqual([], client.stopped)

    def test_cleanup_promotes_capacity_from_complete_active_listing(self) -> None:
        records = [
            {
                "taskId": index,
                "modelId": "owner/old",
                "gpuType": "Iluvatar_bi-100",
                "status": "waiting",
            }
            for index in range(1, 151)
        ]
        client = FakeQueueClient(records)
        pool = ModelHubClientPool(
            [client],  # type: ignore[list-item]
            active_task_cap=100,
            recent_model_reserve_slots=10,
            capacity_state_path=None,
        )
        summary = cleanup_certain_oom_tasks(
            pool,
            FakeDiscovery(
                {"owner/old": 1 * GIB},
                {"owner/old": datetime(2026, 7, 1, tzinfo=timezone.utc)},
            ),  # type: ignore[arg-type]
            reference_time=datetime(2026, 8, 11, tzinfo=timezone.utc),
            log=lambda _message: None,
        )

        self.assertEqual([150], summary["accountCapacityLimits"])
        self.assertEqual([140], summary["oldModelQueueThresholds"])
        self.assertEqual(list(range(141, 151)), [item["taskId"] for item in summary["oldOverflowTasks"]])

    def test_initial_cleanup_stops_old_task_beyond_capacity_minus_ten(self) -> None:
        records = [
            {
                "taskId": index,
                "modelId": "owner/old",
                "gpuType": "Iluvatar_bi-100",
                "status": "waiting",
            }
            for index in range(1, 92)
        ]
        client = FakeQueueClient(records)
        pool = ModelHubClientPool(
            [client],  # type: ignore[list-item]
            active_task_cap=100,
            recent_model_reserve_slots=10,
        )
        summary = cleanup_certain_oom_tasks(
            pool,
            FakeDiscovery(
                {"owner/old": 1 * GIB},
                {"owner/old": datetime(2026, 7, 1, tzinfo=timezone.utc)},
            ),  # type: ignore[arg-type]
            reference_time=datetime(2026, 8, 11, tzinfo=timezone.utc),
            log=lambda _message: None,
        )

        self.assertEqual(1, summary["oldOverflowCount"])
        self.assertEqual(1, summary["cancelledCount"])
        self.assertEqual([90], summary["oldModelQueueThresholds"])
        self.assertEqual([[91]], client.stopped)

    def test_scheduled_cleanup_uses_capacity_minus_five(self) -> None:
        records = [
            {
                "taskId": index,
                "modelId": "owner/old",
                "gpuType": "Iluvatar_bi-100",
                "status": "waiting",
            }
            for index in range(1, 97)
        ]
        client = FakeQueueClient(records)
        pool = ModelHubClientPool(
            [client],  # type: ignore[list-item]
            active_task_cap=100,
            recent_model_reserve_slots=10,
        )
        summary = cleanup_certain_oom_tasks(
            pool,
            FakeDiscovery(
                {"owner/old": 1 * GIB},
                {"owner/old": datetime(2026, 7, 1, tzinfo=timezone.utc)},
            ),  # type: ignore[arg-type]
            age_reserved_slots=5,
            reference_time=datetime(2026, 8, 11, tzinfo=timezone.utc),
            log=lambda _message: None,
        )

        self.assertEqual([95], summary["oldModelQueueThresholds"])
        self.assertEqual([96], [item["taskId"] for item in summary["oldOverflowTasks"]])
        self.assertEqual([[96]], client.stopped)

    def test_oom_is_removed_before_recalculating_old_overflow_positions(self) -> None:
        records = [
            {
                "taskId": index,
                "modelId": "owner/large" if index == 1 else "owner/old",
                "gpuType": "Iluvatar_bi-100",
                "status": "waiting",
            }
            for index in range(1, 93)
        ]
        client = FakeQueueClient(records)
        pool = ModelHubClientPool(
            [client],  # type: ignore[list-item]
            active_task_cap=100,
            recent_model_reserve_slots=10,
        )
        summary = cleanup_certain_oom_tasks(
            pool,
            FakeDiscovery(
                {"owner/large": 40 * GIB, "owner/old": 1 * GIB},
                {"owner/old": datetime(2026, 7, 1, tzinfo=timezone.utc)},
            ),  # type: ignore[arg-type]
            reference_time=datetime(2026, 8, 11, tzinfo=timezone.utc),
            log=lambda _message: None,
        )

        self.assertEqual(1, summary["certainOomCount"])
        self.assertEqual([92], [item["taskId"] for item in summary["oldOverflowTasks"]])
        self.assertEqual([[1], [92]], client.stopped)

    def test_stop_tasks_uses_documented_put_endpoint_and_integer_ids(self) -> None:
        http = RecordingHttpClient()
        client = ModelHubClient(http_client=http)  # type: ignore[arg-type]

        client.stop_tasks(["12", 12, 13])

        self.assertEqual(
            [
                (
                    "PUT",
                    "/api/async/task/stop-create-contest-task",
                    None,
                    {"taskIds": [12, 13]},
                )
            ],
            http.calls,
        )

    def test_only_exact_size_capacity_failures_are_selected(self) -> None:
        tasks = [
            OwnedTask(0, 1, "owner/too-large", "Iluvatar_bi-100", "waiting"),
            OwnedTask(0, 2, "owner/fits", "Iluvatar_bi-100", "waiting"),
            OwnedTask(0, 3, "owner/unknown-size", "Iluvatar_bi-100", "waiting"),
            OwnedTask(0, 4, "owner/unknown-gpu", "new-gpu", "waiting"),
        ]

        selected, skipped = find_certain_oom_tasks(
            tasks,
            repository_sizes={
                "owner/too-large": 30 * GIB,
                "owner/fits": 20 * GIB,
                "owner/unknown-gpu": 30 * GIB,
            },
        )

        self.assertEqual([1], [item["taskId"] for item in selected])
        self.assertEqual(1, skipped["repositorySizeUnknown"])
        self.assertEqual(1, skipped["gpuCapacityUnknown"])
        self.assertEqual(1, skipped["fitsKnownCapacity"])

    def test_cleanup_stops_only_certain_oom_tasks_on_the_owning_account(self) -> None:
        first = FakeQueueClient(
            [
                {"taskId": 1, "modelId": "owner/large", "gpuType": "Iluvatar_bi-100", "status": "waiting"},
                {"taskId": 2, "modelId": "owner/small", "gpuType": "Iluvatar_bi-100", "status": "waiting"},
            ]
        )
        second = FakeQueueClient(
            [{"taskId": 3, "modelId": "owner/large", "gpuType": "MetaX_c-500", "status": "waiting"}]
        )
        pool = ModelHubClientPool([first, second], active_task_cap=100)  # type: ignore[list-item]

        summary = cleanup_certain_oom_tasks(
            pool,
            FakeDiscovery({"owner/large": 40 * GIB, "owner/small": 20 * GIB}),  # type: ignore[arg-type]
            log=lambda _message: None,
        )

        self.assertEqual(1, summary["certainOomCount"])
        self.assertEqual(1, summary["cancelledCount"])
        self.assertEqual([[1]], first.stopped)
        self.assertEqual([], second.stopped)

    def test_task_that_disappears_during_scan_is_not_stopped(self) -> None:
        client = FakeQueueClient(
            [{"taskId": 1, "modelId": "owner/large", "gpuType": "Iluvatar_bi-100", "status": "waiting"}],
            disappear_on_recheck=True,
        )
        pool = ModelHubClientPool([client], active_task_cap=100)  # type: ignore[list-item]

        summary = cleanup_certain_oom_tasks(
            pool,
            FakeDiscovery({"owner/large": 40 * GIB}),  # type: ignore[arg-type]
            log=lambda _message: None,
        )

        self.assertEqual(1, summary["certainOomCount"])
        self.assertEqual(0, summary["cancelledCount"])
        self.assertEqual(1, summary["noLongerActiveCount"])
        self.assertEqual([], client.stopped)


if __name__ == "__main__":
    unittest.main()