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

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)


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

    provider: str
    model_id: str

    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


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

    role: str
    main: ModelRef
    backups: List[ModelRef]
    results: List[Dict[str, Any]]
    assigned_at: float
    rollback_to: Optional[ModelRef] = 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 = None  # Would be thread-safe in production

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

    def get(self, role: str) -> RoleModelAssignment | None:
        """Retrieve a role assignment by role name."""
        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."""
        assignment = self.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."""
        assignment = self.get(role)
        if assignment is None:
            raise NoAvailableModel(f"No assignment found for role {role}")

        if main_available:
            return assignment.main

        # Try backups, prioritizing verified free models
        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."""
        assignment = self.get(role)
        if assignment is None:
            return None
        return assignment.main

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

    def distinct_models(self) -> set[str]:
        """Return distinct model IDs used across all roles."""
        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