from pydantic import BaseModel, Field
from typing import Dict, List, Optional, Any, Literal
from datetime import datetime
from enum import Enum

# ============= ROBOT COMMAND MODELS =============

class JointUpdate(BaseModel):
    name: str
    value: float
    speed: Optional[float] = None

class RobotCommand(BaseModel):
    timestamp: int
    joints: List[JointUpdate]
    duration: Optional[int] = None
    metadata: Optional[Dict[str, Any]] = None

class CommandSequence(BaseModel):
    id: str
    name: str
    commands: List[RobotCommand]
    loop: Optional[bool] = False
    total_duration: int

# ============= COMMUNICATION MODELS =============

class ControlMessage(BaseModel):
    type: Literal["command", "sequence", "heartbeat", "status_request"]
    timestamp: str
    data: Optional[Dict[str, Any]] = None

class StatusMessage(BaseModel):
    type: Literal["status", "joint_states", "error", "heartbeat_ack", "slave_status"]
    timestamp: str
    slave_id: Optional[str] = None
    data: Optional[Dict[str, Any]] = None

# ============= ROBOT MODELS =============

class JointLimits(BaseModel):
    lower: Optional[float] = None
    upper: Optional[float] = None
    velocity: Optional[float] = None
    effort: Optional[float] = None

class DriverJointState(BaseModel):
    name: str
    servo_id: int
    type: Literal["revolute", "continuous"]
    virtual_value: float
    real_value: Optional[float] = None
    limits: Optional[JointLimits] = None

class ConnectionStatus(BaseModel):
    is_connected: bool
    last_connected: Optional[datetime] = None
    error: Optional[str] = None

# ============= ROBOT MANAGEMENT =============

class Robot(BaseModel):
    id: str
    name: str
    robot_type: str
    created_at: datetime
    joints: List[DriverJointState] = []
    
    # Connection state
    master_connected: bool = False
    master_connection_id: Optional[str] = None
    slave_connections: List[str] = []
    
    # Control state
    last_command_source: Literal["master", "manual", "none"] = "none"
    last_command_time: Optional[datetime] = None

class RobotStatus(BaseModel):
    robot_id: str
    master_connected: bool
    slave_count: int
    last_command_source: str
    last_command_time: Optional[datetime]
    last_seen: datetime

class CreateRobotRequest(BaseModel):
    name: Optional[str] = None
    robot_type: str = "so-arm100"

# ============= DEMO SEQUENCES =============

DEMO_SEQUENCES = [
    CommandSequence(
        id="gentle-wave",
        name="Gentle Wave Pattern",
        total_duration=6000,
        commands=[
            RobotCommand(
                timestamp=0,
                joints=[
                    JointUpdate(name="Rotation", value=-10),
                    JointUpdate(name="Pitch", value=8),
                    JointUpdate(name="Elbow", value=-12)
                ],
                duration=2000
            ),
            RobotCommand(
                timestamp=2000,
                joints=[JointUpdate(name="Wrist_Roll", value=10)],
                duration=1000
            ),
            RobotCommand(
                timestamp=3000,
                joints=[JointUpdate(name="Wrist_Roll", value=-10)],
                duration=1000
            ),
            RobotCommand(
                timestamp=4000,
                joints=[
                    JointUpdate(name="Wrist_Roll", value=0),
                    JointUpdate(name="Rotation", value=0),
                    JointUpdate(name="Pitch", value=0),
                    JointUpdate(name="Elbow", value=0)
                ],
                duration=2000
            )
        ]
    ),
    CommandSequence(
        id="small-scan",
        name="Small Scanning Pattern",
        total_duration=8000,
        commands=[
            RobotCommand(
                timestamp=0,
                joints=[
                    JointUpdate(name="Rotation", value=-15),
                    JointUpdate(name="Pitch", value=10)
                ],
                duration=2000
            ),
            RobotCommand(
                timestamp=2000,
                joints=[JointUpdate(name="Rotation", value=15)],
                duration=3000
            ),
            RobotCommand(
                timestamp=5000,
                joints=[
                    JointUpdate(name="Rotation", value=0),
                    JointUpdate(name="Pitch", value=0)
                ],
                duration=3000
            )
        ]
    ),
    CommandSequence(
        id="tiny-flex",
        name="Tiny Flex Pattern",
        total_duration=8000,
        commands=[
            RobotCommand(
                timestamp=0,
                joints=[
                    JointUpdate(name="Elbow", value=-15),
                    JointUpdate(name="Wrist_Pitch", value=8)
                ],
                duration=2000
            ),
            RobotCommand(
                timestamp=2000,
                joints=[JointUpdate(name="Jaw", value=8)],
                duration=1000
            ),
            RobotCommand(
                timestamp=3000,
                joints=[JointUpdate(name="Elbow", value=-25)],
                duration=2000
            ),
            RobotCommand(
                timestamp=5000,
                joints=[
                    JointUpdate(name="Jaw", value=0),
                    JointUpdate(name="Elbow", value=0),
                    JointUpdate(name="Wrist_Pitch", value=0)
                ],
                duration=3000
            )
        ]
    )
]