"""Role assignment management for Model Intelligence with free-only constraints."""

from __future__ import annotations

import logging
import threading
import time
from dataclasses import dataclass, field
from typing import Any

logger = logging.getLogger(__name__)


@dataclass
class ModelRef:
    """Reference to a model in the registry."""

    provider: str
    model_id: str
    is_free: bool = True
    verified: bool = True

    def __post_init__(self) -> None:
        if not self.provider or not self.model_id:
            raise ValueError("ModelRef must have non-empty provider and model_id")


class NoAvailableModel(Exception):
    """Raised when no suitable model is available for assignment."""

    pass


@dataclass
class RoleModelAssignment:
    """Represents a role assignment with main model and backup models."""

    role: str
    main: ModelRef
    backups: list[ModelRef] = field(default_factory=list)
    results: list[dict[str, Any]] = field(default_factory=list)
    assigned_at: float = 0.0
    rollback_to: ModelRef | None = None

    def __repr__(self) -> str:
        return f"RoleModelAssignment(role={self.role}, main={self.main}, backups={len(self.backups)})"


class RoleAssignmentStore:
    """Thread-safe storage for role assignments and their lifecycle."""

    def __init__(self) -> None:
        self._store: dict[str, RoleModelAssignment] = {}
        self._lock = threading.Lock()

    def assign(
        self,
        role: str,
        main: ModelRef,
        backups: list[ModelRef] | None = None,
        *,
        verified: bool = True,
    ) -> RoleModelAssignment:
        """Assign a role with main model and backups."""
        assignment = RoleModelAssignment(
            role=role,
            main=main,
            backups=list(backups) if backups else [],
            results=[],
            assigned_at=time.time(),
        )
        with self._lock:
            self._store[role] = assignment
        logger.info(
            "Assigned role %r with main=%s and backups=%d",
            role,
            main,
            len(assignment.backups),
        )
        return assignment

    def get(self, role: str) -> RoleModelAssignment | None:
        """Retrieve a role assignment by role name."""
        with self._lock:
            return self._store.get(role)

    def record_result(
        self,
        role: str,
        task_kind: str,
        model_id: str,
        success: bool,
    ) -> None:
        """Record a result for a role assignment."""
        with self._lock:
            assignment = self._store.get(role)
            if assignment is None:
                logger.warning("No assignment found for role %r", role)
                return
            assignment.results.append({
                "task_kind": task_kind,
                "model_id": model_id,
                "success": success,
            })

    def select_for_role(
        self,
        role: str,
        *,
        main_available: bool = True,
    ) -> ModelRef | None:
        """Select the main model for a role, falling back to a verified free backup."""
        with self._lock:
            assignment = self._store.get(role)
            if assignment is None:
                raise NoAvailableModel(f"No assignment found for role {role}")

            if main_available:
                return assignment.main

            # Only a backup that is both verified and explicitly free may stand
            # in for an unavailable main model.
            for backup in assignment.backups:
                if backup.verified and backup.is_free:
                    logger.info("Using backup %s for role %r", backup, role)
                    return backup

        raise NoAvailableModel(f"No available model for role {role}")

    def rollback(self, role: str) -> ModelRef | None:
        """Rollback to the main model for a role."""
        with self._lock:
            assignment = self._store.get(role)
            if assignment is None:
                return None
            return assignment.main

    def list_roles(self) -> list[str]:
        """List all defined roles."""
        with self._lock:
            return list(self._store.keys())

    def distinct_models(self) -> set[str]:
        """Return distinct model IDs used across all roles."""
        with self._lock:
            models = set()
            for assignment in self._store.values():
                models.add(assignment.main.model_id)
                for backup in assignment.backups:
                    models.add(backup.model_id)
            return models