import openai
import os
from typing import Dict, Any
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

class CoderAgent:
    """
    Agent responsible for generating code based on user prompts.
    This agent understands programming requirements and writes initial code.
    """
    
    def __init__(self, model="gpt-3.5-turbo", temperature=0.7):
        """
        Initialize the coder agent with specific model parameters.
        
        Args:
            model: Which OpenAI model to use
            temperature: Controls creativity (0=deterministic, 1=creative)
        """
        self.model = model
        self.temperature = temperature
        self.api_key = os.getenv("OPENAI_API_KEY")
        
        if not self.api_key:
            raise ValueError("OPENAI_API_KEY not found in environment variables")
        
        openai.api_key = self.api_key
    
    def generate_code(self, prompt: str, context: str = "") -> Dict[str, Any]:
        """
        Generate code based on user prompt and optional context.
        
        Args:
            prompt: User's coding requirement (e.g., "Write a function to reverse a string")
            context: Additional context from RAG or previous conversations
            
        Returns:
            Dictionary containing code and metadata
        """
        try:
            # Build the system message
            system_message = """You are an expert Python programmer. 
            Write clean, efficient, and well-documented code.
            Always include function signatures with type hints.
            Provide only the code without explanations unless asked."""
            
            # Build user message with context
            user_message = prompt
            if context:
                user_message = f"Context: {context}\n\nTask: {prompt}"
            
            # Call OpenAI API
            response = openai.ChatCompletion.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": system_message},
                    {"role": "user", "content": user_message}
                ],
                temperature=self.temperature,
                max_tokens=500
            )
            
            # Extract and clean the code
            raw_code = response.choices[0].message.content
            
            # Clean the response (remove markdown code blocks if present)
            if "```python" in raw_code:
                code = raw_code.split("```python")[1].split("```")[0].strip()
            elif "```" in raw_code:
                code = raw_code.split("```")[1].split("```")[0].strip()
            else:
                code = raw_code.strip()
            
            return {
                "status": "success",
                "code": code,
                "raw_response": raw_code,
                "model_used": self.model,
                "tokens_used": response.usage.total_tokens
            }
            
        except Exception as e:
            return {
                "status": "error",
                "error": str(e),
                "code": "",
                "raw_response": ""
            }