"""
Intake Configuration Schema
===========================

Pydantic models for intake.yaml configuration files.

These models define the structure of pre-configuration that users can
provide to skip interview questions and speed up study setup.
"""

from __future__ import annotations

from pathlib import Path
from typing import Optional, List, Literal, Union, Any, Dict
from pydantic import BaseModel, Field, field_validator, model_validator
import yaml


class ObjectiveConfig(BaseModel):
    """Configuration for an optimization objective."""

    goal: Literal["minimize", "maximize"]
    target: str = Field(
        description="What to optimize: mass, displacement, stress, frequency, stiffness, or custom name"
    )
    weight: float = Field(default=1.0, ge=0.0, le=10.0)
    extractor: Optional[str] = Field(
        default=None, description="Custom extractor function name (auto-detected if not specified)"
    )

    @field_validator("target")
    @classmethod
    def validate_target(cls, v: str) -> str:
        """Normalize target names."""
        known_targets = {
            "mass",
            "weight",
            "displacement",
            "deflection",
            "stress",
            "frequency",
            "stiffness",
            "strain_energy",
            "volume",
        }
        normalized = v.lower().strip()
        # Map common aliases
        aliases = {
            "weight": "mass",
            "deflection": "displacement",
        }
        return aliases.get(normalized, normalized)


class ConstraintConfig(BaseModel):
    """Configuration for an optimization constraint."""

    type: str = Field(
        description="Constraint type: max_stress, max_displacement, min_frequency, etc."
    )
    threshold: float
    units: str = ""
    description: Optional[str] = None

    @field_validator("type")
    @classmethod
    def normalize_type(cls, v: str) -> str:
        """Normalize constraint type names."""
        return v.lower().strip().replace(" ", "_")


class DesignVariableConfig(BaseModel):
    """Configuration for a design variable."""

    name: str = Field(description="NX expression name")
    bounds: tuple[float, float] = Field(description="(min, max) bounds")
    units: Optional[str] = None
    description: Optional[str] = None
    step: Optional[float] = Field(default=None, description="Step size for discrete variables")

    @field_validator("bounds")
    @classmethod
    def validate_bounds(cls, v: tuple[float, float]) -> tuple[float, float]:
        """Ensure bounds are valid."""
        if len(v) != 2:
            raise ValueError("Bounds must be a tuple of (min, max)")
        if v[0] >= v[1]:
            raise ValueError(f"Lower bound ({v[0]}) must be less than upper bound ({v[1]})")
        return v

    @property
    def range(self) -> float:
        """Get the range of the design variable."""
        return self.bounds[1] - self.bounds[0]

    @property
    def range_ratio(self) -> float:
        """Get the ratio of upper to lower bound."""
        if self.bounds[0] == 0:
            return float("inf")
        return self.bounds[1] / self.bounds[0]


class BudgetConfig(BaseModel):
    """Configuration for optimization budget."""

    max_trials: int = Field(default=100, ge=1, le=10000)
    timeout_per_trial: int = Field(default=300, ge=10, le=7200, description="Seconds per FEA solve")
    target_runtime: Optional[str] = Field(
        default=None, description="Target total runtime (e.g., '2h', '30m')"
    )

    def get_target_runtime_seconds(self) -> Optional[int]:
        """Parse target_runtime string to seconds."""
        if not self.target_runtime:
            return None

        runtime = self.target_runtime.lower().strip()

        if runtime.endswith("h"):
            return int(float(runtime[:-1]) * 3600)
        elif runtime.endswith("m"):
            return int(float(runtime[:-1]) * 60)
        elif runtime.endswith("s"):
            return int(float(runtime[:-1]))
        else:
            # Assume seconds
            return int(float(runtime))


class AlgorithmConfig(BaseModel):
    """Configuration for optimization algorithm."""

    method: Literal["auto", "TPE", "CMA-ES", "NSGA-II", "random"] = "auto"
    neural_acceleration: bool = Field(
        default=False, description="Enable surrogate model for speedup"
    )
    priority: Literal["speed", "accuracy", "balanced"] = "balanced"
    seed: Optional[int] = Field(default=None, description="Random seed for reproducibility")


class MaterialConfig(BaseModel):
    """Configuration for material properties."""

    name: str
    yield_stress: Optional[float] = Field(default=None, ge=0, description="Yield stress in MPa")
    ultimate_stress: Optional[float] = Field(
        default=None, ge=0, description="Ultimate stress in MPa"
    )
    density: Optional[float] = Field(default=None, ge=0, description="Density in kg/m3")
    youngs_modulus: Optional[float] = Field(
        default=None, ge=0, description="Young's modulus in GPa"
    )
    poissons_ratio: Optional[float] = Field(
        default=None, ge=0, le=0.5, description="Poisson's ratio"
    )


class ObjectivesConfig(BaseModel):
    """Configuration for all objectives."""

    primary: ObjectiveConfig
    secondary: Optional[List[ObjectiveConfig]] = None

    @property
    def is_multi_objective(self) -> bool:
        """Check if this is a multi-objective problem."""
        return self.secondary is not None and len(self.secondary) > 0

    @property
    def all_objectives(self) -> List[ObjectiveConfig]:
        """Get all objectives as a flat list."""
        objectives = [self.primary]
        if self.secondary:
            objectives.extend(self.secondary)
        return objectives


class StudyConfig(BaseModel):
    """Configuration for study metadata."""

    name: Optional[str] = Field(
        default=None, description="Study name (auto-generated from folder if omitted)"
    )
    type: Literal["single_objective", "multi_objective"] = "single_objective"
    description: Optional[str] = None
    tags: Optional[List[str]] = None


class IntakeConfig(BaseModel):
    """
    Complete intake.yaml configuration schema.

    All fields are optional - anything not specified will be asked
    in the interview or auto-detected from introspection.
    """

    study: Optional[StudyConfig] = None
    objectives: Optional[ObjectivesConfig] = None
    constraints: Optional[List[ConstraintConfig]] = None
    design_variables: Optional[List[DesignVariableConfig]] = None
    budget: Optional[BudgetConfig] = None
    algorithm: Optional[AlgorithmConfig] = None
    material: Optional[MaterialConfig] = None
    notes: Optional[str] = None

    @classmethod
    def from_yaml(cls, yaml_path: Union[str, Path]) -> "IntakeConfig":
        """Load configuration from a YAML file."""
        yaml_path = Path(yaml_path)

        if not yaml_path.exists():
            raise FileNotFoundError(f"Intake config not found: {yaml_path}")

        with open(yaml_path, "r", encoding="utf-8") as f:
            data = yaml.safe_load(f)

        if data is None:
            return cls()

        return cls.model_validate(data)

    @classmethod
    def from_yaml_safe(cls, yaml_path: Union[str, Path]) -> Optional["IntakeConfig"]:
        """Load configuration from YAML, returning None if file doesn't exist."""
        yaml_path = Path(yaml_path)

        if not yaml_path.exists():
            return None

        try:
            return cls.from_yaml(yaml_path)
        except Exception:
            return None

    def to_yaml(self, yaml_path: Union[str, Path]) -> None:
        """Save configuration to a YAML file."""
        yaml_path = Path(yaml_path)

        data = self.model_dump(exclude_none=True)

        with open(yaml_path, "w", encoding="utf-8") as f:
            yaml.dump(data, f, default_flow_style=False, sort_keys=False)

    def get_value(self, key: str) -> Optional[Any]:
        """
        Get a configuration value by dot-notation key.

        Examples:
            config.get_value("study.name")
            config.get_value("budget.max_trials")
            config.get_value("objectives.primary.goal")
        """
        parts = key.split(".")
        value: Any = self

        for part in parts:
            if value is None:
                return None
            if hasattr(value, part):
                value = getattr(value, part)
            elif isinstance(value, dict):
                value = value.get(part)
            else:
                return None

        return value

    def is_complete(self) -> bool:
        """Check if all required configuration is provided."""
        return (
            self.objectives is not None
            and self.design_variables is not None
            and len(self.design_variables) > 0
        )

    def get_missing_fields(self) -> List[str]:
        """Get list of fields that still need to be configured."""
        missing = []

        if self.objectives is None:
            missing.append("objectives")
        if self.design_variables is None or len(self.design_variables) == 0:
            missing.append("design_variables")
        if self.constraints is None:
            missing.append("constraints (recommended)")
        if self.budget is None:
            missing.append("budget")

        return missing

    @model_validator(mode="after")
    def validate_consistency(self) -> "IntakeConfig":
        """Validate consistency between configuration sections."""
        # Check study type matches objectives
        if self.study and self.objectives:
            is_multi = self.objectives.is_multi_objective
            declared_multi = self.study.type == "multi_objective"

            if is_multi and not declared_multi:
                # Auto-correct study type
                self.study.type = "multi_objective"

        return self


# Common material presets
MATERIAL_PRESETS: Dict[str, MaterialConfig] = {
    "aluminum_6061_t6": MaterialConfig(
        name="Aluminum 6061-T6",
        yield_stress=276,
        ultimate_stress=310,
        density=2700,
        youngs_modulus=68.9,
        poissons_ratio=0.33,
    ),
    "aluminum_7075_t6": MaterialConfig(
        name="Aluminum 7075-T6",
        yield_stress=503,
        ultimate_stress=572,
        density=2810,
        youngs_modulus=71.7,
        poissons_ratio=0.33,
    ),
    "steel_a36": MaterialConfig(
        name="Steel A36",
        yield_stress=250,
        ultimate_stress=400,
        density=7850,
        youngs_modulus=200,
        poissons_ratio=0.26,
    ),
    "stainless_304": MaterialConfig(
        name="Stainless Steel 304",
        yield_stress=215,
        ultimate_stress=505,
        density=8000,
        youngs_modulus=193,
        poissons_ratio=0.29,
    ),
    "titanium_6al4v": MaterialConfig(
        name="Titanium Ti-6Al-4V",
        yield_stress=880,
        ultimate_stress=950,
        density=4430,
        youngs_modulus=113.8,
        poissons_ratio=0.342,
    ),
}


def get_material_preset(name: str) -> Optional[MaterialConfig]:
    """
    Get a material preset by name (fuzzy matching).

    Examples:
        get_material_preset("6061")  # Returns aluminum_6061_t6
        get_material_preset("steel")  # Returns steel_a36
    """
    name_lower = name.lower().replace("-", "_").replace(" ", "_")

    # Direct match
    if name_lower in MATERIAL_PRESETS:
        return MATERIAL_PRESETS[name_lower]

    # Partial match
    for key, material in MATERIAL_PRESETS.items():
        if name_lower in key or name_lower in material.name.lower():
            return material

    return None