"""
LLM Interface for interacting with AI models
"""
from typing import Dict, List, Optional, Any
import aiohttp
import json
import os

class LLMInterface:
    def __init__(self, api_key: str, model: str = "gpt-4"):
        self.api_key = api_key
        self.model = model
        self.base_url = "https://api.openai.com/v1"
        self.session = None
        
    async def __aenter__(self):
        self.session = aiohttp.ClientSession()
        return self
        
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.session:
            await self.session.close()
            
    async def analyze_request(self, request: str, context: Dict) -> Dict[str, Any]:
        """Analyze user request and determine required actions"""
        prompt = self._build_analysis_prompt(request, context)
        response = await self._call_api(prompt)
        return self._parse_analysis_response(response)
    
    async def generate_solution(self, task: str, context: Dict) -> Dict[str, Any]:
        """Generate solution for a specific task"""
        prompt = self._build_solution_prompt(task, context)
        response = await self._call_api(prompt)
        return self._parse_solution_response(response)
    
    async def _call_api(self, prompt: str) -> Dict:
        """Make API call to LLM service"""
        if not self.session:
            self.session = aiohttp.ClientSession()
            
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        data = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": "You are an expert programming assistant."},
                {"role": "user", "content": prompt}
            ],
            "temperature": 0.7
        }
        
        async with self.session.post(
            f"{self.base_url}/chat/completions",
            headers=headers,
            json=data
        ) as response:
            return await response.json()
            
    def _build_analysis_prompt(self, request: str, context: Dict) -> str:
        """Build prompt for request analysis"""
        return f"""
        Given the following user request and workspace context, analyze what needs to be done:
        
        REQUEST:
        {request}
        
        CONTEXT:
        {json.dumps(context, indent=2)}
        
        Please provide:
        1. Required actions
        2. Files that need to be modified
        3. Dependencies that need to be added
        4. Potential risks or concerns
        """
        
    def _build_solution_prompt(self, task: str, context: Dict) -> str:
        """Build prompt for solution generation"""
        return f"""
        Generate a solution for the following task with given context:
        
        TASK:
        {task}
        
        CONTEXT:
        {json.dumps(context, indent=2)}
        
        Please provide:
        1. Detailed implementation steps
        2. Code changes required
        3. Test cases needed
        4. Validation steps
        """
        
    def _parse_analysis_response(self, response: Dict) -> Dict[str, Any]:
        """Parse LLM response for analysis"""
        # Extract the content from the response
        content = response["choices"][0]["message"]["content"]
        
        # Parse the structured response
        # This is a simple implementation - could be enhanced with better parsing
        sections = content.split("\n\n")
        parsed = {
            "actions": [],
            "files": [],
            "dependencies": [],
            "risks": []
        }
        
        current_section = None
        for line in content.split("\n"):
            if "Required actions:" in line:
                current_section = "actions"
            elif "Files that need to be modified:" in line:
                current_section = "files"
            elif "Dependencies that need to be added:" in line:
                current_section = "dependencies"
            elif "Potential risks or concerns:" in line:
                current_section = "risks"
            elif current_section and line.strip():
                parsed[current_section].append(line.strip())
                
        return parsed
        
    def _parse_solution_response(self, response: Dict) -> Dict[str, Any]:
        """Parse LLM response for solution"""
        content = response["choices"][0]["message"]["content"]
        
        # Parse the structured response
        parsed = {
            "steps": [],
            "code_changes": [],
            "tests": [],
            "validation": []
        }
        
        current_section = None
        for line in content.split("\n"):
            if "Detailed implementation steps:" in line:
                current_section = "steps"
            elif "Code changes required:" in line:
                current_section = "code_changes"
            elif "Test cases needed:" in line:
                current_section = "tests"
            elif "Validation steps:" in line:
                current_section = "validation"
            elif current_section and line.strip():
                parsed[current_section].append(line.strip())
                
        return parsed
