import openai
from typing import Dict, Any, Optional

class AIProjectAssistant:
    """Wrapper class for OpenAI API interactions focused on Data Science and AI projects."""
    
    def __init__(self, api_key: str):
        """Initialize the AI assistant with OpenAI API key."""
        openai.api_key = api_key
        self.default_model = "gpt-4"
        self.default_temperature = 0.7
        
    def send_prompt(self, 
                    prompt: str, 
                    model: Optional[str] = None,
                    temperature: Optional[float] = None) -> Dict[str, Any]:
        """
        Send a prompt to OpenAI API and get the response.
        
        Args:
            prompt (str): The user's input prompt
            model (str, optional): OpenAI model to use. Defaults to gpt-4
            temperature (float, optional): Response creativity (0-1). Defaults to 0.7
            
        Returns:
            Dict[str, Any]: OpenAI API response
        """
        try:
            response = openai.ChatCompletion.create(
                model=model or self.default_model,
                messages=[
                    {"role": "user", "content": prompt}
                ],
                temperature=temperature or self.default_temperature
            )
            return {
                "success": True,
                "content": response.choices[0].message.content,
                "tokens_used": response.usage.total_tokens
            }
        except Exception as e:
            return {
                "success": False,
                "error": str(e)
            }
    
    def brainstorm_project(self, topic: str) -> Dict[str, Any]:
        """
        Generate AI project ideas based on a topic.
        
        Args:
            topic (str): The topic or field of interest
            
        Returns:
            Dict[str, Any]: Project suggestions and implementation details
        """
        prompt = f"""
        Help me brainstorm an AI/Data Science project related to {topic}. Please provide:
        1. Project title
        2. Problem statement
        3. Suggested approach
        4. Required technologies/libraries
        5. Potential challenges
        6. Expected outcomes
        """
        return self.send_prompt(prompt)
    
    def get_code_suggestion(self, description: str) -> Dict[str, Any]:
        """
        Get code suggestions for implementing specific functionality.
        
        Args:
            description (str): Description of the desired functionality
            
        Returns:
            Dict[str, Any]: Code suggestions and explanations
        """
        prompt = f"""
        Please provide Python code suggestions for the following functionality:
        {description}
        
        Include:
        1. Code implementation
        2. Required imports
        3. Brief explanation of the approach
        4. Any potential improvements
        """
        return self.send_prompt(prompt)