agents.py

"""
Multi-agent system for Blender code generation
- Planning Agent: Interprets and breaks down user prompts
- Code Generation Agent: Generates Blender Python code
"""
import json
from typing import Dict, List, Optional
from .claude_client import ClaudeClient

class PlanningAgent:
    """Agent that interprets user prompts and breaks them down into actionable steps"""
    
    SYSTEM_PROMPT = """Blender planning agent. Return JSON only:
{"understanding": "brief", "steps": [{"step_number": 1, "description": "action"}], "complexity": "simple|medium|complex", "estimated_operations": []}"""

    def __init__(self, client: ClaudeClient, use_fast_model: bool = True):
        self.client = client
        # Use faster model for planning if available
        if use_fast_model and "haiku" not in client.model.lower():
            # Create a fast client for planning
            self.fast_client = ClaudeClient(client.api_key, "claude-3-5-haiku-20241022")
        else:
            self.fast_client = client
    
    async def plan(self, user_prompt: str) -> Dict:
        """Generate a plan from user prompt (optimized for speed)"""
        # Check if it's a simple request that can skip planning - be more aggressive
        simple_keywords = ["cube", "sphere", "cylinder", "torus", "plane", "simple", "add", "create", "make"]
        is_simple = any(keyword in user_prompt.lower() for keyword in simple_keywords) and len(user_prompt.split()) < 20
        
        if is_simple:
            # Skip planning for very simple requests
            return {
                "understanding": user_prompt,
                "steps": [{"step_number": 1, "description": user_prompt, "blender_operations": [], "details": ""}],
                "complexity": "simple",
                "estimated_operations": []
            }
        
        messages = [{"role": "user", "content": user_prompt}]
        
        # Use faster model and fewer tokens for planning
        response = await self.fast_client.chat(
            messages=messages,
            system=self.SYSTEM_PROMPT,
            max_tokens=512  # Reduced for faster responses
        )
        
        # Try to extract JSON from response
        try:
            # Find JSON in response (might have markdown code blocks)
            if "```json" in response:
                json_start = response.find("```json") + 7
                json_end = response.find("```", json_start)
                response = response[json_start:json_end].strip()
            elif "```" in response:
                json_start = response.find("```") + 3
                json_end = response.find("```", json_start)
                response = response[json_start:json_end].strip()
            
            plan = json.loads(response)
            return plan
        except json.JSONDecodeError:
            # Fallback if JSON parsing fails
            return {
                "understanding": response,
                "steps": [{"step_number": 1, "description": response, "blender_operations": [], "details": ""}],
                "complexity": "medium",
                "estimated_operations": []
            }


class CodeGenerationAgent:
    """Agent that generates Blender Python code based on plans"""
    
    SYSTEM_PROMPT = """Generate Blender Python code. Rules: import bpy, use bpy.ops, executable code only. No markdown."""

    def __init__(self, client: ClaudeClient):
        self.client = client
    
    async def generate_code(self, plan: Dict, user_prompt: str, context: Optional[str] = None) -> str:
        """Generate Blender Python code from plan (optimized for speed)"""
        # Ultra-simplified prompts for speed
        if plan.get("complexity") == "simple" and len(plan.get("steps", [])) == 1:
            # For simple requests, minimal prompt
            prompt = f"Blender code: {user_prompt}"
        else:
            # For complex requests, minimal plan reference
            plan_summary = plan.get("understanding", user_prompt)
            prompt = f"Blender code: {user_prompt}"
        
        messages = [{"role": "user", "content": prompt}]
        
        code = await self.client.chat(
            messages=messages,
            system=self.SYSTEM_PROMPT,
            max_tokens=2048  # Reduced for faster generation - most code is < 500 tokens
        )
        
        # Clean up code - remove markdown if present
        if "```python" in code:
            code_start = code.find("```python") + 9
            code_end = code.find("```", code_start)
            code = code[code_start:code_end].strip()
        elif "```" in code:
            code_start = code.find("```") + 3
            code_end = code.find("```", code_start)
            code = code[code_start:code_end].strip()
        
        return code
    
    def get_blender_context(self) -> str:
        """Get current Blender scene context as string"""
        try:
            context_info = []
            context_info.append(f"Active Object: {bpy.context.active_object.name if bpy.context.active_object else 'None'}")
            context_info.append(f"Selected Objects: {[obj.name for obj in bpy.context.selected_objects]}")
            context_info.append(f"Scene Objects: {len(bpy.data.objects)} objects")
            context_info.append(f"Materials: {len(bpy.data.materials)} materials")
            return "\n".join(context_info)
        except Exception as e:
            return f"Error getting context: {str(e)}"