"""
Pydantic request models for API endpoints.
"""

from pydantic import BaseModel, Field, ConfigDict
from typing import List, Dict, Any, Optional, Literal, Union


# ============================================================
# LLM / AI-assisted design models
# ============================================================

class LLMProviderConfig(BaseModel):
    """Configuration for a structuring LLM provider (OpenAI or Ollama)."""
    provider: Literal["openai", "ollama"] = Field(
        ..., description="Provider identifier"
    )
    model: str = Field(
        ..., description="Model name, e.g. 'gpt-4o', 'gpt-4.1', 'llama3.2'"
    )
    api_key: Optional[str] = Field(
        None, description="API key (required for openai; omit for ollama)"
    )
    base_url: Optional[str] = Field(
        None, description="Base URL override (for custom Ollama installs)"
    )


class EdisonConfig(BaseModel):
    """Optional Edison Scientific literature search configuration."""
    api_key: Optional[str] = Field(
        None,
        description=(
            "Edison platform API key. If None, the SDK uses the EDISON_API_KEY env var."
        ),
    )
    job_type: Literal["literature", "literature_high", "precedent"] = Field(
        "literature",
        description=(
            "'literature' — standard PaperQA3 search; "
            "'literature_high' — high-reasoning mode; "
            "'precedent' — HasAnyone-style precedent search"
        ),
    )
    timeout_secs: Optional[int] = Field(
        None,
        ge=60,
        le=3600,
        description=(
            "Maximum seconds to wait for the Edison response (default 1200 = 20 min). "
            "If the task is not complete by this deadline, the search is abandoned and "
            "the structuring model proceeds without literature context."
        ),
    )
    force_refresh: bool = Field(
        False,
        description=(
            "If True, ignore any cached result and re-submit a fresh search to Edison, "
            "even if a completed result is already stored for this query."
        ),
    )


class SuggestEffectsRequest(BaseModel):
    """Request body for POST /api/v1/llm/suggest-effects/{session_id}."""
    structuring_provider: LLMProviderConfig = Field(
        ..., description="LLM used to extract the structured effects list"
    )
    edison_config: Optional[EdisonConfig] = Field(
        None,
        description=(
            "If provided, Edison Scientific is queried first for grounded "
            "literature context; its cited answer is then fed to the structuring model."
        ),
    )
    system_context: str = Field(
        ...,
        description=(
            "Free-text description of the experimental system and optimization target, "
            "e.g. 'Fischer-Tropsch synthesis over supported metal catalysts, "
            "maximizing C5+ selectivity at 250 °C and 20 bar.'"
        ),
    )


# ============================================================
# Variable Models
# ============================================================

class AddRealVariableRequest(BaseModel):
    """Request to add a real-valued variable."""
    name: str = Field(..., description="Variable name")
    type: Literal["real"] = Field(default="real", description="Variable type")
    min: float = Field(..., description="Minimum value")
    max: float = Field(..., description="Maximum value")
    unit: Optional[str] = Field(None, description="Unit of measurement")
    description: Optional[str] = Field(None, description="Variable description")
    
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "name": "temperature",
                "type": "real",
                "min": 300,
                "max": 500,
                "unit": "°C",
                "description": "Reaction temperature"
            }
        }
    )


class AddIntegerVariableRequest(BaseModel):
    """Request to add an integer variable."""
    name: str = Field(..., description="Variable name")
    type: Literal["integer"] = Field(default="integer", description="Variable type")
    min: int = Field(..., description="Minimum value")
    max: int = Field(..., description="Maximum value")
    unit: Optional[str] = Field(None, description="Unit of measurement")
    description: Optional[str] = Field(None, description="Variable description")
    
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "name": "batch_size",
                "type": "integer",
                "min": 1,
                "max": 10,
                "unit": "batches",
                "description": "Number of batches"
            }
        }
    )


class AddCategoricalVariableRequest(BaseModel):
    """Request to add a categorical variable."""
    name: str = Field(..., description="Variable name")
    type: Literal["categorical"] = Field(default="categorical", description="Variable type")
    categories: List[str] = Field(..., description="List of category values")
    unit: Optional[str] = Field(None, description="Unit of measurement")
    description: Optional[str] = Field(None, description="Variable description")
    
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "name": "catalyst",
                "type": "categorical",
                "categories": ["A", "B", "C"],
                "description": "Catalyst type"
            }
        }
    )


class AddDiscreteVariableRequest(BaseModel):
    """Request to add a discrete numerical variable."""
    name: str = Field(..., description="Variable name")
    type: Literal["discrete"] = Field(default="discrete", description="Variable type")
    allowed_values: List[float] = Field(..., min_length=2, description="List of allowed numeric values (at least 2)")
    unit: Optional[str] = Field(None, description="Unit of measurement")
    description: Optional[str] = Field(None, description="Variable description")

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "name": "SAR",
                "type": "discrete",
                "allowed_values": [80, 280],
                "unit": "-",
                "description": "Silicon-to-aluminium ratio (only synthesizable at specific values)"
            }
        }
    )


# Union type for any variable request
AddVariableRequest = Union[
    AddRealVariableRequest,
    AddIntegerVariableRequest,
    AddCategoricalVariableRequest,
    AddDiscreteVariableRequest
]


# ============================================================
# Experiment Models
# ============================================================

class AddExperimentRequest(BaseModel):
    """Request to add a single experiment."""
    inputs: Dict[str, Union[float, int, str]] = Field(..., description="Variable values")
    output: Optional[float] = Field(None, description="Target/output value")
    noise: Optional[float] = Field(None, description="Measurement uncertainty")
    iteration: Optional[int] = Field(None, description="Iteration number (auto-assigned if None)")
    reason: Optional[str] = Field(None, description="Reason for this experiment")
    
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "inputs": {"temperature": 350, "catalyst": "A"},
                "output": 0.85,
                "noise": 0.02,
                "iteration": 1,
                "reason": "Initial Design"
            }
        }
    )


class StageExperimentRequest(BaseModel):
    """Request to stage an experiment for later execution."""
    inputs: Dict[str, Union[float, int, str]] = Field(..., description="Variable values")
    reason: Optional[str] = Field(None, description="Reason for this experiment (e.g., acquisition strategy)")
    
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "inputs": {"temperature": 375.2, "catalyst": "B"},
                "reason": "qEI"
            }
        }
    )


class StageExperimentsBatchRequest(BaseModel):
    """Request to stage multiple experiments at once."""
    experiments: List[Dict[str, Union[float, int, str]]] = Field(..., description="List of experiment inputs")
    reason: Optional[str] = Field(None, description="Reason for these experiments")
    
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "experiments": [
                    {"temperature": 375.2, "catalyst": "B"},
                    {"temperature": 412.8, "catalyst": "A"}
                ],
                "reason": "qEI batch"
            }
        }
    )


class CompleteStagedExperimentsRequest(BaseModel):
    """Request to complete staged experiments with outputs."""
    outputs: List[float] = Field(..., description="Output values for staged experiments (same order)")
    noises: Optional[List[float]] = Field(None, description="Measurement uncertainties (optional)")
    iteration: Optional[int] = Field(None, description="Iteration number (auto-assigned if None)")
    reason: Optional[str] = Field(None, description="Reason (uses staged reason if not provided)")
    
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "outputs": [0.87, 0.92],
                "noises": [0.02, 0.03],
                "iteration": 5,
                "reason": "qEI"
            }
        }
    )


class AddExperimentsBatchRequest(BaseModel):
    """Request to add multiple experiments."""
    experiments: List[AddExperimentRequest] = Field(..., description="List of experiments")
    
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "experiments": [
                    {"inputs": {"temperature": 350, "catalyst": "A"}, "output": 0.85},
                    {"inputs": {"temperature": 400, "catalyst": "B"}, "output": 0.92}
                ]
            }
        }
    )


# ============================================================
# Model Training Models
# ============================================================

class TrainModelRequest(BaseModel):
    """Request to train a surrogate model."""
    backend: Literal["sklearn", "botorch"] = Field(default="sklearn", description="Modeling backend")
    kernel: str = Field(default="Matern", description="Kernel type (RBF, Matern, RationalQuadratic for sklearn; RBF, Matern, IBNN for botorch)")
    kernel_params: Optional[Dict[str, Any]] = Field(None, description="Kernel-specific parameters")
    input_transform: Optional[str] = Field(None, description="Input transformation (Normalize, Standardize, etc.)")
    output_transform: Optional[str] = Field(None, description="Output transformation (Standardize, etc.)")
    calibration_enabled: bool = Field(default=False, description="Enable uncertainty calibration")
    
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "backend": "sklearn",
                "kernel": "Matern",
                "kernel_params": {"nu": 2.5}
            }
        }
    )


# ============================================================
# Acquisition Models
# ============================================================

class AcquisitionRequest(BaseModel):
    """Request to suggest next experiments."""
    strategy: str = Field(default="EI", description="Acquisition strategy (EI, PI, UCB, qEI, qUCB, qNIPV)")
    goal: Literal["maximize", "minimize"] = Field(default="maximize", description="Optimization goal")
    n_suggestions: int = Field(default=1, ge=1, le=10, description="Number of suggestions (batch size)")
    xi: Optional[float] = Field(default=0.01, description="Exploration parameter for EI/PI")
    kappa: Optional[float] = Field(default=2.0, description="Exploration parameter for UCB")
    
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "strategy": "EI",
                "goal": "maximize",
                "n_suggestions": 1,
                "xi": 0.01
            }
        }
    )


class FindOptimumRequest(BaseModel):
    """Request to find model's predicted optimum."""
    goal: Literal["maximize", "minimize"] = Field(default="maximize", description="Optimization goal")
    
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "goal": "maximize"
            }
        }
    )


# ============================================================
# Initial Design (DoE) Models
# ============================================================

class InitialDesignRequest(BaseModel):
    """Request for generating initial experimental design.

    Space-filling methods (random, lhs, sobol, halton, hammersly) require n_points.
    Classical RSM methods (full_factorial, fractional_factorial, ccd, box_behnken)
    determine run count from design structure; n_points is ignored.
    """
    method: Literal[
        "random", "lhs", "sobol", "halton", "hammersly",
        "full_factorial", "fractional_factorial", "ccd", "box_behnken",
        "plackett_burman", "gsd"
    ] = Field(default="lhs", description="Sampling method")
    n_points: Optional[int] = Field(
        None, ge=1, le=1000,
        description="Number of points (required for space-filling, ignored for classical designs)"
    )
    random_seed: Optional[int] = Field(None, description="Random seed for reproducibility")
    lhs_criterion: str = Field(
        default="maximin",
        pattern="^(maximin|correlation|ratio)$",
        description="Criterion for LHS method"
    )
    # Classical design parameters
    n_levels: int = Field(default=2, ge=2, le=5, description="Levels per factor (full factorial)")
    n_center: int = Field(default=1, ge=0, le=10, description="Center point replicates")
    generators: Optional[str] = Field(None, description="Fractional factorial generator string")
    ccd_alpha: Literal["orthogonal", "rotatable"] = Field(
        default="orthogonal", description="CCD alpha type"
    )
    ccd_face: Literal["circumscribed", "inscribed", "faced"] = Field(
        default="circumscribed", description="CCD face type"
    )
    gsd_reduction: int = Field(
        default=2, ge=2, le=10,
        description="GSD reduction factor (larger = fewer runs)"
    )

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "method": "lhs",
                "n_points": 10,
                "random_seed": 42,
                "lhs_criterion": "maximin"
            }
        }
    )


# ============================================================
# Optimal Design Models
# ============================================================

class OptimalDesignInfoRequest(BaseModel):
    """Request to preview optimal design model terms and recommended run count.

    Performs a dry-run inspection without running the exchange algorithm.
    Specify either model_type (shortcut) or effects (explicit), not both.
    """
    model_type: Optional[Literal["linear", "interaction", "quadratic"]] = Field(
        None, description="Shortcut model type"
    )
    effects: Optional[List[str]] = Field(
        None,
        description=(
            "Explicit effect strings using variable names. "
            "Main effects: 'Temperature'; interactions: 'Temperature*Pressure'; "
            "quadratic: 'Temperature**2'. Intercept is added automatically."
        ),
    )

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "model_type": "quadratic"
            }
        }
    )


class OptimalDesignRequest(BaseModel):
    """Request to generate a statistically optimal experimental design.

    Specify either model_type (shortcut) or effects (explicit), not both.
    Specify either n_points (absolute) or p_multiplier (relative to model columns), not both.
    """
    model_type: Optional[Literal["linear", "interaction", "quadratic"]] = Field(
        None, description="Shortcut model type"
    )
    effects: Optional[List[str]] = Field(
        None,
        description=(
            "Explicit effect strings using variable names. "
            "Main effects: 'Temperature'; interactions: 'Temperature*Pressure'; "
            "quadratic: 'Temperature**2'. Intercept is added automatically."
        ),
    )
    n_points: Optional[int] = Field(
        None, ge=1, le=10000,
        description="Absolute number of experimental runs"
    )
    p_multiplier: Optional[float] = Field(
        None, ge=1.0, le=10.0,
        description="Run count as multiple of model columns p (e.g. 2.0 → 2p runs)"
    )
    criterion: Literal["D", "A", "I"] = Field(
        default="D",
        description="Optimality criterion: D (parameter estimation), A (min avg variance), I (min prediction variance)"
    )
    algorithm: Literal[
        "sequential", "simple_exchange", "fedorov", "modified_fedorov", "detmax"
    ] = Field(default="fedorov", description="Exchange algorithm")
    n_levels: int = Field(
        default=5, ge=2, le=20,
        description="Candidate grid levels per continuous variable"
    )
    max_iter: int = Field(
        default=200, ge=10, le=10000,
        description="Maximum exchange iterations"
    )
    random_seed: Optional[int] = Field(None, description="Random seed for reproducibility")

    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "model_type": "quadratic",
                "p_multiplier": 2.0,
                "criterion": "D",
                "algorithm": "fedorov"
            }
        }
    )


# ============================================================
# Prediction Models
# ============================================================

class PredictionRequest(BaseModel):
    """Request to make predictions at new points."""
    inputs: List[Dict[str, Union[float, int, str]]] = Field(..., description="Input points for prediction")
    
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "inputs": [
                    {"temperature": 375, "catalyst": "A"},
                    {"temperature": 425, "catalyst": "B"}
                ]
            }
        }
    )


# ============================================================
# Audit Log & Session Management Models
# ============================================================

class UpdateMetadataRequest(BaseModel):
    """Request to update session metadata."""
    name: Optional[str] = Field(None, description="Session name")
    description: Optional[str] = Field(None, description="Session description")
    tags: Optional[List[str]] = Field(None, description="Session tags")
    
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "name": "Catalyst_Screening_Nov2025",
                "description": "Pt/Pd ratio optimization for CO2 reduction",
                "tags": ["catalyst", "CO2", "electrochemistry"]
            }
        }
    )


class LockDecisionRequest(BaseModel):
    """Request to lock in a decision to the audit log."""
    lock_type: Literal["data", "model", "acquisition"] = Field(..., description="Type of decision to lock")
    notes: Optional[str] = Field(None, description="Optional notes about this decision")
    
    # For acquisition lock
    strategy: Optional[str] = Field(None, description="Acquisition strategy (required for acquisition lock)")
    parameters: Optional[Dict[str, Any]] = Field(None, description="Acquisition parameters (required for acquisition lock)")
    suggestions: Optional[List[Dict[str, Any]]] = Field(None, description="Suggested experiments (required for acquisition lock)")
    
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "lock_type": "model",
                "notes": "Best cross-validation performance: R²=0.93"
            }
        }
    )


# ============================================================
# Session Lock Models
# ============================================================

class SessionLockRequest(BaseModel):
    """Request to lock a session for programmatic control."""
    locked_by: str = Field(..., description="Identifier of the client locking the session")
    client_id: Optional[str] = Field(None, description="Optional unique client identifier")
    
    model_config = ConfigDict(
        json_schema_extra={
            "example": {
                "locked_by": "Reactor Controller v1.2",
                "client_id": "lab-3-workstation"
            }
        }
    )
