data_structures.py

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

# --- ORCHESTRATION ---
class UserGoal(BaseModel):
    user_id: str
    project_id: str
    goal_text: str
    roblox_user_id: Optional[str] = "0"
    image_base64: Optional[str] = Field(None, description="Base64 encoded context image.")

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

# --- FEEDBACK LOOP ---
class SceneObjectState(BaseModel):
    """Represents a snapshot of an object currently in Roblox."""
    id: str # The Name in Roblox
    class_name: str
    position: Dict[str, float]
    size: Dict[str, float]
    parent_name: str

class ProjectFeedback(BaseModel):
    """Data sent FROM the Plugin TO the Orchestrator."""
    project_id: str
    message: str
    scene_snapshot: List[SceneObjectState] = []
    task_id_completed: Optional[str] = None
    status: Literal["SUCCESS", "FAILED"] = "SUCCESS"

# --- MATTER FORMAT ---
class Vector3(BaseModel):
    x: float = 0.0
    y: float = 0.0
    z: float = 0.0

class BuildPayload(BaseModel):
    id: str
    class_name: str = "Part"
    position: Vector3
    rotation: Vector3 = Vector3()
    size: Vector3
    color_hex: str = "#A3A2A5"
    anchored: bool = True
    transparency: float = 0.0
    material: str = "Plastic"
    mesh_id: Optional[str] = None
    texture_id: Optional[str] = None

class CSGPayload(BaseModel):
    operation: Literal["Union", "Subtract"]
    target_id: str
    tool_ids: List[str]

class Attribute(BaseModel):
    name: str
    type: Literal["String", "Number", "Boolean", "Color3"]
    value: Any

class ScriptPayload(BaseModel):
    name: str
    class_name: Literal["Script", "LocalScript", "ModuleScript"]
    parent: str 
    attributes: List[Attribute] = [] 
    source_code: str 

class RiggingCommand(BaseModel):
    part0_id: str
    part1_id: str
    joint_type: Literal["Weld", "Motor6D", "HingeConstraint", "BallSocketConstraint"]
    position: Vector3

# --- PROJECT FLOW ---
class TaskNode(BaseModel):
    task_id: str = Field(default_factory=lambda: "pending_id")
    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
    is_finished: bool = False
    next_task: Optional[TaskNode] = None

# --- COMMUNICATION ---
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