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

# --- 0. ORCHESTRATION MODELS ---
class UserGoal(BaseModel):
    user_id: str
    project_id: str
    goal_text: str
    image_base64: Optional[str] = None

class RecallRequest(BaseModel):
    project_id: str
    user_feedback: Optional[str] = "What is the next step?"

class ResetRequest(BaseModel):
    junior_instance_id: str
    supervisor_instance_id: str
    reason: str

# --- 1. THE MATTER FORMAT (Physical World) ---
class Vector3(BaseModel):
    x: float = 0.0
    y: float = 0.0
    z: float = 0.0

class MatterObject(BaseModel):
    """Represents a primitive part or mesh."""
    id: str
    class_name: str = "Part" # Part, MeshPart, TrussPart, WedgePart
    position: Vector3
    rotation: Vector3
    size: Vector3
    color_hex: str = "#A3A2A5"
    anchored: bool = True
    transparency: float = 0.0
    material: str = "Plastic"
    mesh_id: Optional[str] = None # For custom meshes
    texture_id: Optional[str] = None # rbxassetid://...

class ToolboxAsset(BaseModel):
    """Represents a pre-made model from Roblox Toolbox."""
    id: str
    asset_id: str # The numeric Roblox Asset ID (e.g. "123456")
    position: Vector3
    rotation: Vector3
    scale: float = 1.0

class RiggingCommand(BaseModel):
    """Instructions to rig/connect parts."""
    part0_id: str
    part1_id: str
    joint_type: Literal["Weld", "Motor6D", "HingeConstraint", "BallSocketConstraint"]
    position: Vector3 # World position where the joint sits

# --- 2. THE PROJECT FLOW FORMAT ---
class TaskNode(BaseModel):
    task_id: str
    title: str
    role_required: Literal[
        "PROJECT_MANAGER", 
        "SUPERVISOR_WORLD", 
        "SUPERVISOR_SCRIPT", 
        "3D_ARTIST", 
        "CONCEPT_ARTIST", 
        "SCRIPTING_ENGINEER", 
        "ERRANDS"
    ]
    instruction: str
    status: Literal["PENDING", "IN_PROGRESS", "COMPLETED", "FAILED"] = "PENDING"
    assigned_instance_id: Optional[str] = None
    subtasks: List['TaskNode'] = [] 

class ProjectFlow(BaseModel):
    project_id: str
    game_concept_summary: str
    root_tasks: List[TaskNode]

# --- 3. COMMUNICATION PROTOCOLS ---
class AgentMessage(BaseModel):
    id: str
    from_agent: str
    to_agents: List[str]
    role: str
    content: str
    image_base64: Optional[str] = None

class GenerationRequest(BaseModel):
    agent_type: str
    agent_instance_id: str
    conversation_partner_id: str 
    prompt: str
    image_base64: Optional[str] = None

class ImageGenerationRequest(BaseModel):
    agent_instance_id: str
    prompt: str