from __future__ import annotations

import threading
from pathlib import Path
from typing import Any

from common import read_jsonl, update_jsonl, utc_now


DEFAULT_SUBMISSION_EXCLUSIONS_PATH = Path(".modelhub_state/submission_exclusions.jsonl")


def exclusion_key(model_id: str, target_gpu: str) -> str:
    return f"{model_id}|{target_gpu}"


class SubmissionExclusionStore:
    """Persistent model/GPU exclusions for non-retryable platform rejections."""

    def __init__(self, path: Path | str = DEFAULT_SUBMISSION_EXCLUSIONS_PATH) -> None:
        self.path = Path(path)
        self._lock = threading.Lock()
        self._blocked = {
            exclusion_key(str(record.get("modelId")), str(record.get("targetGpu")))
            for record in read_jsonl(self.path)
            if record.get("modelId") and record.get("targetGpu")
        }

    def is_blocked(self, model_id: str, target_gpu: str) -> bool:
        with self._lock:
            return exclusion_key(model_id, target_gpu) in self._blocked

    def block(self, model_id: str, target_gpu: str, *, reason: str) -> None:
        if not model_id or not target_gpu:
            return
        key = exclusion_key(model_id, target_gpu)
        now = utc_now().isoformat()

        def update(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
            existing = next(
                (
                    record
                    for record in records
                    if exclusion_key(str(record.get("modelId")), str(record.get("targetGpu"))) == key
                ),
                None,
            )
            if existing is None:
                records.append(
                    {
                        "modelId": model_id,
                        "targetGpu": target_gpu,
                        "reason": reason,
                        "firstSeenAt": now,
                        "lastSeenAt": now,
                        "occurrences": 1,
                    }
                )
            else:
                existing["reason"] = reason
                existing["lastSeenAt"] = now
                existing["occurrences"] = max(1, int(existing.get("occurrences") or 1)) + 1
            return records

        update_jsonl(self.path, update)
        with self._lock:
            self._blocked.add(key)