import os
import sys
from typing import List, Dict
import anthropic
from dotenv import load_dotenv

class ClaudeChat:
    def __init__(self, max_history: int = 10):
        """
        Initialize Claude chat client with history management
        
        Args:
            max_history (int): Maximum number of messages to keep in history
        """
        # Load environment variables
        load_dotenv()

        # Retrieve API key
        self.api_key = os.getenv('ClaudeAI_API_KEY')
        if not self.api_key:
            print("Error: ClaudeAI_API_KEY not found in environment variables")
            sys.exit(1)

        # Initialize client
        self.client = anthropic.Anthropic(api_key=self.api_key)

        # Conversation history
        self.history: List[Dict] = []
        self.max_history = max_history
        
        # Model selection
        self.model = "claude-3-5-sonnet-latest"
        
        # System prompt
        self.system_prompt = """You are an expert computer science and data science educator. Your task is to create clear, engaging, and pedagogically sound lecture slides. Follow these guidelines:

1. Structure:
- Start with learning objectives
- Present concepts progressively from simple to complex
- Include concrete examples and visualizations
- End with summary/key takeaways

2. Content:
- Use clear, concise language
- Include practical examples and code snippets
- Add visualizations wherever possible
- Highlight key concepts and terminology

3. Format:
- Use consistent slide templates
- Include code chunks with proper formatting
- Hide implementation details when appropriate
- Use color coding for important concepts

4. Examples:
- Start with small, tractable examples
- Show step-by-step problem solving
- Include both successful and error cases
- Connect to real-world applications

5. Technical details:
- Use #| echo: false for hidden code chunks
- use R code to generate plots/graphs where necessary, R generate image files to be embedded in the slides
- R code for generating plots/graphs should be embedded in slides but hidden
- Use latex for tables when necessary
- use "$" for inline latex math/expression and "$$" for block latex
- Include proper figure dimensions
- Ensure code outputs are readable
- Test all code examples before including
- output quarto beamer slide deck.

Provide clear, educational content suitable for undergraduate/graduate level computer science and data science courses."""
        self.add_to_history("system", self.system_prompt)

    def add_to_history(self, role: str, content: str):
        """
        Add message to conversation history
        
        Args:
            role (str): 'user' or 'assistant'
            content (str): Message content
        """
        self.history.append({"role": role, "content": content})
        
        # Trim history if it exceeds max_history
        if len(self.history) > self.max_history:
            self.history.pop(0)

    def chat(self):
        """
        Interactive chat loop with Claude
        """
        print(f"🤖 Claude Chat (Model: {self.model})")
        print("Type 'exit' or 'quit' to end the conversation.")
        print("Type 'clear' to reset conversation history.\n")

        while True:
            try:
                # Get user input
                user_input = input("You: ").strip()

                # Handle special commands
                if user_input.lower() in ['exit', 'quit']:
                    print("Goodbye!")
                    break
                
                if user_input.lower() == 'clear':
                    self.history.clear()
                    print("Conversation history cleared.")
                    continue

                # Add user message to history
                self.add_to_history("user", user_input)

                # Prepare messages for API call
                messages = [
                    {"role": msg["role"], "content": msg["content"]} 
                    for msg in self.history
                ]

                # Create message with Claude
                response = self.client.messages.create(
                    model=self.model,
                    max_tokens=14096,
                    messages=messages
                )

                # Extract and print Claude's response
                claude_response = response.content[0].text
                print(f"\nClaude: {claude_response}\n")

                # Add assistant response to history
                self.add_to_history("assistant", claude_response)

            except KeyboardInterrupt:
                print("\nChat interrupted. Type 'exit' to quit.")
            except Exception as e:
                print(f"An error occurred: {e}")

def main():
    chat_client = ClaudeChat()
    chat_client.chat()

if __name__ == "__main__":
    main()