"""
ControlPlan v0.1: Declarative HIL control specification.

This model defines a JSON-serializable spec that describes HIL behavior
using high-level tasks (loops, playback profiles) and actions (set, add, if).
The spec is compiled deterministically into Python HIL logic.

Design goals:
- LLM-friendly: structured output, no free-form Python
- Safe: expressions are parsed with AST, only safe operations allowed
- Expressive: loops, conditionals, profiles (Gaussian, ramp, etc.)
"""

from __future__ import annotations

from typing import Dict, List, Literal, Optional, Tuple, Union

from pydantic import BaseModel, ConfigDict, Field, field_validator


# =============================================================================
# Action types
# =============================================================================

class SetAction(BaseModel):
    """Set a variable to an expression result: var = expr"""
    model_config = ConfigDict(extra="forbid")

    set: Tuple[str, str] = Field(
        description="[variable_name, expression_string]"
    )


class AddAction(BaseModel):
    """Add expression result to a variable: var += expr"""
    model_config = ConfigDict(extra="forbid")

    add: Tuple[str, str] = Field(
        description="[variable_name, expression_string]"
    )


class IfAction(BaseModel):
    """
    Conditional action: if condition then actions [else actions].

    The condition is an expression string that evaluates to a boolean.
    """
    model_config = ConfigDict(extra="forbid")

    # Using "if_" to avoid Python keyword conflict, aliased to "if" in JSON
    if_: str = Field(alias="if", description="Condition expression string")
    then: List["Action"] = Field(description="Actions to execute if condition is true")
    else_: Optional[List["Action"]] = Field(
        default=None,
        alias="else",
        description="Actions to execute if condition is false"
    )


# Union of all action types
Action = Union[SetAction, AddAction, IfAction]

# Enable forward reference resolution
IfAction.model_rebuild()


# =============================================================================
# Profile types (for playback tasks)
# =============================================================================

class GaussianProfile(BaseModel):
    """Gaussian noise profile for playback tasks."""
    model_config = ConfigDict(extra="forbid")

    kind: Literal["gaussian"] = "gaussian"
    height: float = Field(description="Base/center value")
    mean: float = Field(default=0.0, description="Mean of Gaussian noise")
    std: float = Field(gt=0, description="Standard deviation of Gaussian noise")
    entries: int = Field(gt=0, description="Number of entries in one cycle")


class RampProfile(BaseModel):
    """Linear ramp profile for playback tasks."""
    model_config = ConfigDict(extra="forbid")

    kind: Literal["ramp"] = "ramp"
    start: float = Field(description="Start value")
    end: float = Field(description="End value")
    entries: int = Field(gt=0, description="Number of entries in one cycle")


class StepProfile(BaseModel):
    """Step function profile for playback tasks."""
    model_config = ConfigDict(extra="forbid")

    kind: Literal["step"] = "step"
    values: List[float] = Field(min_length=1, description="Values to cycle through")


Profile = Union[GaussianProfile, RampProfile, StepProfile]


# =============================================================================
# Task types
# =============================================================================

class LoopTask(BaseModel):
    """
    A loop task that executes actions repeatedly at a fixed interval.

    This is the main mechanism for implementing control logic:
    - Read inputs (from physical_values)
    - Compute outputs (via expressions)
    - Write outputs (via set/add actions)
    """
    model_config = ConfigDict(extra="forbid")

    type: Literal["loop"] = "loop"
    name: str = Field(description="Task name for debugging/logging")
    dt_s: float = Field(gt=0, description="Loop interval in seconds")
    actions: List[Action] = Field(description="Actions to execute each iteration")


class PlaybackTask(BaseModel):
    """
    A playback task that outputs a profile (Gaussian, ramp, step) over time.

    Use for generating test signals, disturbances, or time-varying setpoints.
    """
    model_config = ConfigDict(extra="forbid")

    type: Literal["playback"] = "playback"
    name: str = Field(description="Task name for debugging/logging")
    dt_s: float = Field(gt=0, description="Interval between profile entries")
    target: str = Field(description="Variable name to write profile values to")
    profile: Profile = Field(description="Profile definition")
    repeat: bool = Field(default=True, description="Whether to repeat the profile")


Task = Union[LoopTask, PlaybackTask]


# =============================================================================
# Top-level structures
# =============================================================================

class ControlPlanHIL(BaseModel):
    """
    HIL control plan: defines initialization and tasks for one HIL.

    Structure:
    - name: must match hils[].name in configuration.json
    - warmup_s: optional delay before starting tasks
    - init: initial values for state variables
    - params: optional constants for use in expressions
    - tasks: list of loop/playback tasks (run in separate threads if >1)
    """
    model_config = ConfigDict(extra="forbid")

    name: str = Field(description="HIL name (must match config)")
    warmup_s: Optional[float] = Field(
        default=None,
        ge=0,
        description="Warmup delay in seconds before starting tasks"
    )
    init: Dict[str, Union[float, int, bool]] = Field(
        description="Initial values for physical_values keys"
    )
    params: Optional[Dict[str, Union[float, int, bool]]] = Field(
        default=None,
        description="Constants available in expressions"
    )
    tasks: List[Task] = Field(
        min_length=1,
        description="Tasks to run (loop or playback)"
    )

    @field_validator("init")
    @classmethod
    def validate_init_keys(cls, v: Dict[str, Union[float, int, bool]]) -> Dict[str, Union[float, int, bool]]:
        """Ensure init keys are valid Python identifiers."""
        for key in v.keys():
            if not key.isidentifier():
                raise ValueError(f"Invalid init key (not a valid identifier): {key}")
        return v

    @field_validator("params")
    @classmethod
    def validate_params_keys(cls, v: Optional[Dict[str, Union[float, int, bool]]]) -> Optional[Dict[str, Union[float, int, bool]]]:
        """Ensure params keys are valid Python identifiers."""
        if v is None:
            return v
        for key in v.keys():
            if not key.isidentifier():
                raise ValueError(f"Invalid params key (not a valid identifier): {key}")
        return v


class ControlPlan(BaseModel):
    """
    Top-level ControlPlan specification.

    Usage:
        control_plan.json -> compile_control_plan.py -> hil_*.py

    The generated HIL logic follows the ICS-SimLab contract:
        def logic(physical_values):
            # init phase
            # warmup sleep
            # while True: run tasks
    """
    model_config = ConfigDict(extra="forbid")

    version: Literal["v0.1"] = Field(
        default="v0.1",
        description="Schema version"
    )
    hils: List[ControlPlanHIL] = Field(
        min_length=1,
        description="HIL control specifications"
    )


def get_control_plan_json_schema() -> dict:
    """Return JSON Schema for ControlPlan, suitable for LLM structured output."""
    return ControlPlan.model_json_schema()