"""
Pydantic models for simulation configuration.
"""
from typing import List, Optional
from pydantic import BaseModel, Field, field_validator


class ServerConfig(BaseModel):
    """Configuration for a single server in the network."""
    id: str = Field(..., description="Unique server identifier")
    service_rate: float = Field(..., gt=0, description="Service rate μ (services per time unit)")
    routing_probability: float = Field(..., ge=0, le=1, description="Probability q of routing to this server")

    model_config = {
        "json_schema_extra": {
            "example": {
                "id": "server_1",
                "service_rate": 0.00833,
                "routing_probability": 0.5
            }
        }
    }


class SimulationConfigModel(BaseModel):
    """
    Complete configuration for a queueing network simulation.

    This model is used for API requests and validation.
    """
    # Network parameters
    arrival_rate: float = Field(..., gt=0, description="External arrival rate λ (arrivals per time unit)")
    coordinator_service_rate: float = Field(..., gt=0, description="Coordinator service rate μc")
    coordinator_exit_probability: float = Field(..., ge=0, le=1, description="Probability p of exiting after coordinator")
    servers: List[ServerConfig] = Field(..., min_length=1, description="List of server configurations")

    # Simulation control parameters
    warmup_time: float = Field(10000.0, ge=0, description="Warmup period duration (discard initial statistics)")
    simulation_time: float = Field(100000.0, gt=0, description="Total simulation duration")
    random_seed: Optional[int] = Field(None, description="Random seed for reproducibility")

    # Output preferences
    collect_time_series: bool = Field(True, description="Collect time series data (customers over time)")
    time_series_interval: float = Field(100.0, gt=0, description="Time interval for time series sampling")
    collect_histograms: bool = Field(True, description="Collect processing time histograms")
    histogram_bins: int = Field(50, gt=0, description="Number of bins for histograms")

    @field_validator('servers')
    @classmethod
    def validate_probability_conservation(cls, servers: List[ServerConfig], info) -> List[ServerConfig]:
        """Ensure routing probabilities + exit probability sum to 1.0."""
        # We need access to coordinator_exit_probability, which is in info.data
        if 'coordinator_exit_probability' in info.data:
            exit_prob = info.data['coordinator_exit_probability']
            routing_sum = sum(s.routing_probability for s in servers)
            total = exit_prob + routing_sum

            if not (0.99 <= total <= 1.01):  # Allow small floating point errors
                raise ValueError(
                    f"Exit probability ({exit_prob}) + sum of routing probabilities ({routing_sum}) "
                    f"must equal 1.0, got {total}"
                )

        return servers

    @field_validator('simulation_time')
    @classmethod
    def validate_simulation_time(cls, simulation_time: float, info) -> float:
        """Ensure simulation time is greater than warmup time."""
        if 'warmup_time' in info.data:
            warmup_time = info.data['warmup_time']
            if simulation_time <= warmup_time:
                raise ValueError(
                    f"Simulation time ({simulation_time}) must be greater than warmup time ({warmup_time})"
                )
        return simulation_time

    model_config = {
        "json_schema_extra": {
            "example": {
                "arrival_rate": 0.008,
                "coordinator_service_rate": 0.1,
                "coordinator_exit_probability": 0.5,
                "servers": [
                    {
                        "id": "server_1",
                        "service_rate": 0.00833,
                        "routing_probability": 0.5
                    }
                ],
                "warmup_time": 10000.0,
                "simulation_time": 100000.0,
                "random_seed": 42
            }
        }
    }

    def to_simulation_config(self):
        """Convert to internal SimulationConfig for the simulator."""
        from src.core.simulation import SimulationConfig

        return SimulationConfig(
            arrival_rate=self.arrival_rate,
            coordinator_service_rate=self.coordinator_service_rate,
            coordinator_exit_probability=self.coordinator_exit_probability,
            server_service_rates=[s.service_rate for s in self.servers],
            server_routing_probs=[s.routing_probability for s in self.servers],
            warmup_time=self.warmup_time,
            simulation_time=self.simulation_time,
            random_seed=self.random_seed
        )