from __future__ import annotations

import sys
import tempfile
import threading
import unittest
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timezone
from pathlib import Path


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 common import read_json, read_jsonl  # noqa: E402
from main import build_parser, make_run_dir, run_submission, submit_candidate  # noqa: E402
from modelhub_client import ModelHubAPIError, ModelHubClientPool, is_duplicate_submission_error  # noqa: E402
from models import HFModelSummary, ModelInspection  # noqa: E402
from outcome_tracker import OutcomeTracker  # noqa: E402
from submission_claims import SubmissionClaimStore, candidate_key  # noqa: E402


class FakeClient:
    def __init__(self, active_count: int = 0, *, reject_capacity: bool = False) -> None:
        self.active_count = active_count
        self.reject_capacity = reject_capacity
        self.submitted: list[dict] = []
        self.count_calls = 0
        self.search_calls = 0
        self._lock = threading.Lock()

    def count_active_tasks(self, **_kwargs) -> int:  # noqa: ANN003
        with self._lock:
            self.count_calls += 1
            return self.active_count

    def add_task(self, payload: dict) -> dict:
        if self.reject_capacity:
            raise ModelHubAPIError("当前等待中或运行中的异步模型验证任务数量已达上限")
        with self._lock:
            self.submitted.append(payload)
            return {"code": 0, "data": {"id": len(self.submitted)}}

    def search_by_model_id(self, _model_id: str) -> dict:
        with self._lock:
            self.search_calls += 1
        return {"code": 0, "data": {"verifyResult": {}}}

    def list_tasks(self, **_kwargs) -> list[dict]:  # noqa: ANN003
        return []

    def list_tasks_page(self, **_kwargs) -> dict:  # noqa: ANN003
        return {"code": 0, "data": {"records": [], "pages": 0}}

    def find_recent_task_id(self, *_args, **_kwargs):  # noqa: ANN002, ANN003
        return None


class FakeDiscovery:
    def __init__(self, count: int) -> None:
        self.models = [
            HFModelSummary(
                repo_id=f"owner/model-{index}",
                downloads=100,
                last_modified=datetime(2026, 1, 1, tzinfo=timezone.utc),
                pipeline_tag="text-generation",
            )
            for index in range(count)
        ]

    def list_recent_models(self, **_kwargs):  # noqa: ANN003
        return self.models

    @staticmethod
    def inspect_model(model: HFModelSummary) -> ModelInspection:
        return ModelInspection(repo_id=model.repo_id, weight_files=["model.safetensors"])


class DuplicateThenSuccessClient:
    def __init__(self) -> None:
        self.calls = 0
        self.active = 0

    @staticmethod
    def begin_cycle() -> None:
        return None

    def available_submit_slots(self) -> int:
        return max(0, 2 - self.active)

    @staticmethod
    def list_tasks_page(**_kwargs) -> dict:  # noqa: ANN003
        return {"code": 0, "data": {"records": [], "pages": 0}}

    @staticmethod
    def processed_gpus_for_model(_model_id: str) -> set[str]:
        return set()

    def add_task(self, _payload: dict) -> dict:
        self.calls += 1
        if self.calls <= 2:
            raise ModelHubAPIError("模型正在验证中，请勿重复提交")
        self.active += 1
        return {"code": 0, "data": {"id": f"task-{self.calls}"}}


class DynamicCapacityClient:
    def __init__(self, *, active: int, limit: int, token: str = "dynamic-test") -> None:
        self.active = active
        self.limit = limit
        self.token = token

    def count_active_tasks(self, *, max_count: int, **_kwargs) -> int:  # noqa: ANN003
        return min(self.active, max_count)

    def add_task(self, _payload: dict) -> dict:
        if self.active >= self.limit:
            raise ModelHubAPIError("当前等待中或运行中的异步模型验证任务数量已达上限")
        self.active += 1
        return {"code": 0, "data": {"id": self.active}}


class AutoStrategyClient:
    def __init__(self, available: int) -> None:
        self.available = available
        self.calls = 0

    @staticmethod
    def begin_cycle() -> None:
        return None

    def available_submit_slots(self) -> int:
        return self.available

    @staticmethod
    def list_tasks(**_kwargs) -> list[dict]:  # noqa: ANN003
        return []

    @staticmethod
    def list_tasks_page(**_kwargs) -> dict:  # noqa: ANN003
        return {"code": 0, "data": {"records": [], "pages": 0}}

    @staticmethod
    def processed_gpus_for_model(_model_id: str) -> set[str]:
        return set()

    def add_task(self, _payload: dict) -> dict:
        self.calls += 1
        if self.calls == 1:
            raise ModelHubAPIError("模型正在验证中，请勿重复提交")
        self.available -= 1
        return {"code": 0, "data": {"id": f"strategy-{self.calls}"}}


def make_candidate(index: int) -> dict:
    return {
        "repoId": f"owner/model-{index}",
        "modelAddress": f"https://modelscope.cn/models/owner/model-{index}",
        "targetGpu": "NVIDIA-A100",
    }


class ClientPoolConcurrencyTests(unittest.TestCase):
    def test_adaptive_strategy_counts_only_platform_accepted_tasks(self) -> None:
        with tempfile.TemporaryDirectory() as temporary_dir:
            root = Path(temporary_dir)
            args = build_parser().parse_args(
                [
                    "--task-types",
                    "text-generation",
                    "--limit",
                    "4",
                    "--max-scan-models",
                    "4",
                    "--skip-outcome-sync",
                    "--skip-history-archive",
                ]
            )
            args.runs_dir = str(root / "runs")
            args.ledger_path = str(root / "ledger.jsonl")
            args.outcomes_path = str(root / "outcomes.jsonl")
            args.claims_path = str(root / "claims.jsonl")
            args.history_archive_path = str(root / "history.jsonl")
            args.gpu_strategy_state_path = str(root / "strategy.json")
            client = AutoStrategyClient(available=3)

            summary = run_submission(
                args,
                now=datetime(2026, 1, 1, 12, tzinfo=timezone.utc),
                hf_discovery=FakeDiscovery(4),  # type: ignore[arg-type]
                modelhub_client=client,  # type: ignore[arg-type]
            )

            state = read_json(root / "strategy.json")
            self.assertEqual(3, summary["submittedCount"])
            self.assertEqual(1, summary["duplicateCount"])
            self.assertEqual(3, state["acceptedSinceRefresh"])
            self.assertEqual(3, sum(state["acceptedByCategory"].values()))

    def test_duplicate_batch_is_replaced_until_available_slots_are_filled(self) -> None:
        with tempfile.TemporaryDirectory() as temporary_dir:
            root = Path(temporary_dir)
            args = build_parser().parse_args(
                [
                    "--gpus",
                    "Iluvatar_bi-150",
                    "--task-types",
                    "text-generation",
                    "--limit",
                    "4",
                    "--max-scan-models",
                    "4",
                    "--skip-outcome-sync",
                    "--skip-history-archive",
                ]
            )
            args.runs_dir = str(root / "runs")
            args.ledger_path = str(root / "ledger.jsonl")
            args.outcomes_path = str(root / "outcomes.jsonl")
            args.claims_path = str(root / "claims.jsonl")
            args.history_archive_path = str(root / "history.jsonl")
            client = DuplicateThenSuccessClient()

            summary = run_submission(
                args,
                now=datetime(2026, 1, 1, 12, tzinfo=timezone.utc),
                hf_discovery=FakeDiscovery(4),  # type: ignore[arg-type]
                modelhub_client=client,  # type: ignore[arg-type]
            )

            self.assertEqual(2, summary["targetSubmitCount"])
            self.assertEqual(4, summary["plannedSubmitCount"])
            self.assertEqual(2, summary["duplicateCount"])
            self.assertEqual(2, summary["submittedCount"])
            self.assertEqual(0, client.available_submit_slots())

    def test_platform_duplicate_is_classified_as_skipped_candidate(self) -> None:
        error = ModelHubAPIError("模型正在验证中，请勿重复提交")
        self.assertTrue(is_duplicate_submission_error(error))

        class DuplicateClient:
            @staticmethod
            def add_task(_payload):  # noqa: ANN001
                raise error

        candidate = {
            **make_candidate(1),
            "taskType": "text-generation",
            "framework": "vllm",
            "configParams": "framework: vllm",
        }
        result = submit_candidate(candidate, DuplicateClient())  # type: ignore[arg-type]
        self.assertEqual("duplicate", result["outcome"])

    def test_every_account_reaches_capacity_under_load(self) -> None:
        clients = [FakeClient() for _ in range(12)]
        pool = ModelHubClientPool(
            clients,  # type: ignore[arg-type]
            active_task_cap=10,
            active_counts_ttl=60,
            instance_id="full-load-test",
        )

        with ThreadPoolExecutor(max_workers=24) as executor:
            list(executor.map(lambda index: pool.add_task({"index": index}), range(120)))

        self.assertEqual([10] * 12, [len(client.submitted) for client in clients])
        self.assertEqual([10] * 12, pool.active_task_counts())
        self.assertEqual(0, pool.available_submit_slots())

    def test_concurrent_submissions_reserve_and_balance_accounts(self) -> None:
        clients = [FakeClient(), FakeClient()]
        pool = ModelHubClientPool(
            clients,  # type: ignore[arg-type]
            active_task_cap=4,
            active_counts_ttl=60,
            instance_id="balance-test",
        )

        with ThreadPoolExecutor(max_workers=8) as executor:
            results = list(executor.map(lambda index: pool.add_task({"index": index}), range(8)))

        self.assertEqual(8, len(results))
        self.assertEqual([4, 4], [len(client.submitted) for client in clients])
        self.assertEqual([4, 4], pool.active_task_counts())
        self.assertEqual(0, pool.available_submit_slots())
        # A lagging remote count must not erase successful local reservations.
        pool._refresh_active_counts(force=True)
        self.assertEqual([4, 4], pool.active_task_counts())

    def test_capacity_rejection_falls_through_to_another_account(self) -> None:
        full_elsewhere = FakeClient(active_count=0, reject_capacity=True)
        available = FakeClient(active_count=1)
        pool = ModelHubClientPool(
            [full_elsewhere, available],  # type: ignore[arg-type]
            active_task_cap=2,
            instance_id="fallback-test",
        )

        response = pool.add_task({"model": "x"})

        self.assertEqual(1, response["data"]["id"])
        self.assertEqual(0, len(full_elsewhere.submitted))
        self.assertEqual(1, len(available.submitted))

    def test_verification_cache_is_reset_between_cycles(self) -> None:
        client = FakeClient()
        pool = ModelHubClientPool([client], active_task_cap=2)  # type: ignore[arg-type]
        pool.search_by_model_id("owner/model")
        pool.search_by_model_id("owner/model")
        self.assertEqual(1, client.search_calls)
        pool.begin_cycle()
        pool.search_by_model_id("owner/model")
        self.assertEqual(2, client.search_calls)

    def test_capacity_probe_discovers_a_higher_dynamic_account_limit(self) -> None:
        client = DynamicCapacityClient(active=2, limit=3)
        pool = ModelHubClientPool(
            [client],  # type: ignore[arg-type]
            active_task_cap=2,
            active_counts_ttl=60,
            capacity_probe_interval_cycles=3,
            instance_id="capacity-growth-test",
        )

        pool.configure_capacity_probe(2)
        self.assertEqual(0, pool.available_submit_slots())
        pool.configure_capacity_probe(3)
        self.assertEqual(1, pool.available_submit_slots())
        pool.add_task({"model": "probe"})

        self.assertEqual([3], pool.account_capacity_limits())
        self.assertEqual([3], pool.active_task_counts())
        self.assertEqual(0, pool.available_submit_slots())

    def test_rejected_capacity_probe_keeps_known_limit_and_enters_cooldown(self) -> None:
        client = DynamicCapacityClient(active=2, limit=2)
        pool = ModelHubClientPool(
            [client],  # type: ignore[arg-type]
            active_task_cap=2,
            active_counts_ttl=60,
            capacity_probe_interval_cycles=3,
            capacity_probe_cooldown_cycles=3,
            instance_id="capacity-rejection-test",
        )

        pool.configure_capacity_probe(3)
        self.assertEqual(1, pool.available_submit_slots())
        with self.assertRaises(ModelHubAPIError):
            pool.add_task({"model": "probe"})

        self.assertEqual([2], pool.account_capacity_limits())
        self.assertEqual(0, pool.available_submit_slots())


class ProcessCoordinationTests(unittest.TestCase):
    def test_duplicate_claims_are_retained_so_next_batch_moves_forward(self) -> None:
        with tempfile.TemporaryDirectory() as temporary_dir:
            path = Path(temporary_dir) / "claims.jsonl"
            candidates = [make_candidate(index) for index in range(4)]
            store = SubmissionClaimStore(path, owner_id="worker-a")

            duplicates = store.claim(candidates, limit=2)
            store.mark_submitted(duplicates)
            replacements = store.claim(candidates, limit=2)

            self.assertEqual(
                {candidate_key(candidate) for candidate in candidates[2:]},
                {candidate_key(candidate) for candidate in replacements},
            )

    def test_concurrent_claim_stores_select_disjoint_candidates(self) -> None:
        with tempfile.TemporaryDirectory() as temporary_dir:
            path = Path(temporary_dir) / "claims.jsonl"
            candidates = [make_candidate(index) for index in range(6)]
            first = SubmissionClaimStore(path, owner_id="worker-a")
            second = SubmissionClaimStore(path, owner_id="worker-b")

            with ThreadPoolExecutor(max_workers=2) as executor:
                first_future = executor.submit(first.claim, candidates, limit=3)
                second_future = executor.submit(second.claim, candidates, limit=3)
                first_claims = first_future.result()
                second_claims = second_future.result()

            first_keys = {candidate_key(candidate) for candidate in first_claims}
            second_keys = {candidate_key(candidate) for candidate in second_claims}
            self.assertEqual(3, len(first_claims))
            self.assertEqual(3, len(second_claims))
            self.assertTrue(first_keys.isdisjoint(second_keys))

    def test_outcome_saves_merge_instead_of_overwriting_other_worker(self) -> None:
        with tempfile.TemporaryDirectory() as temporary_dir:
            path = Path(temporary_dir) / "outcomes.jsonl"
            first = OutcomeTracker(path)
            second = OutcomeTracker(path)
            first.record_submission("model-a", "gpu", "vllm", "text-generation", "task-a", "2026-01-01T00:00:00+00:00")
            second.record_submission("model-b", "gpu", "vllm", "text-generation", "task-b", "2026-01-01T00:00:01+00:00")

            first.save()
            second.save()

            self.assertEqual({"task-a", "task-b"}, {row["taskId"] for row in read_jsonl(path)})

    def test_run_directories_are_allocated_atomically(self) -> None:
        with tempfile.TemporaryDirectory() as temporary_dir:
            base = Path(temporary_dir) / "runs"
            now = datetime(2026, 1, 1, tzinfo=timezone.utc)
            with ThreadPoolExecutor(max_workers=8) as executor:
                paths = list(executor.map(lambda _index: make_run_dir(base, now), range(8)))

            self.assertEqual(8, len(set(paths)))
            self.assertTrue(all(path.is_dir() for path in paths))


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