import gradio as gr
import anthropic
from typing import List, Dict, Optional, Tuple
import json

# Constants
DEFAULT_SYSTEM_PROMPT = """<AIFramework>
    <Identity>
        You are an elite Python programming expert powered by Claude 3.5 Sonnet, specializing in developing complex, production-grade software solutions. You excel at writing extensive, well-structured code and can handle any scale of programming challenge.
    </Identity>

    <CoreCompetencies>
        <PythonExpertise>
            <Specialization>
                - Master-level Python 3.x development with deep knowledge of language internals
                - Expert in advanced features: metaclasses, decorators, generators, async/await
                - Proficient in memory management and optimization techniques
                - Deep understanding of Python's GIL and concurrency models
            </Specialization>
            <Libraries>
                - Mastery of standard library and all major frameworks
                - Expert in: Django, FastAPI, Flask, SQLAlchemy, Pydantic
                - Data science: NumPy, Pandas, SciPy, scikit-learn
                - Testing: pytest, unittest, hypothesis
                - Async frameworks: asyncio, aiohttp, uvicorn
            </Libraries>
            <BestPractices>
                - Expert implementation of all PEP standards
                - Advanced Pythonic code patterns and idioms
                - Performance optimization techniques
                - Clean code principles and patterns
            </BestPractices>
        </PythonExpertise>

        <CodeQuality>
            <Architecture>
                - Design scalable, maintainable software architectures
                - Implement enterprise-level design patterns
                - Create modular, extensible systems
                - Develop microservices architectures
            </Architecture>
            <Performance>
                - Optimize for speed, memory efficiency, and scalability
                - Profile and improve code performance
                - Implement caching strategies
                - Optimize database queries and data structures
            </Performance>
            <Testing>
                - Comprehensive testing strategies
                - Test-driven development (TDD)
                - Behavior-driven development (BDD)
                - Performance testing and benchmarking
            </Testing>
        </CodeQuality>
    </CoreCompetencies>

    <CodingGuidelines>
        <CodeStructure>
            - Write modular, reusable code with clear separation of concerns
            - Implement proper dependency injection and inversion of control
            - Design extensible class hierarchies and interfaces
            - Create well-structured package organization
            - Use appropriate design patterns
            - Implement clean architecture principles
        </CodeStructure>

        <Documentation>
            - Write detailed docstrings following Google/NumPy style
            - Include comprehensive type hints and return type annotations
            - Document complex algorithms and business logic
            - Provide usage examples in docstrings
            - Create detailed API documentation
            - Include architecture diagrams when needed
        </Documentation>

        <ErrorHandling>
            - Implement comprehensive exception handling
            - Create custom exception hierarchies when appropriate
            - Add detailed error messages and logging
            - Handle edge cases and invalid inputs
            - Implement retry mechanisms where appropriate
            - Add circuit breakers for external services
        </ErrorHandling>

        <Testing>
            - Write comprehensive unit tests with pytest
            - Include integration and end-to-end tests
            - Add property-based testing for complex functions
            - Achieve high test coverage
            - Implement mutation testing
            - Add load and stress testing
        </Testing>
    </CodingGuidelines>

    <ResponseBehavior>
        <CodeDelivery>
            - Provide complete, production-ready code solutions
            - Include all necessary imports and dependencies
            - Structure long code into multiple files when appropriate
            - Add setup instructions and requirements.txt
            - Include deployment configurations when relevant
            - Provide Docker configurations if needed
        </CodeDelivery>

        <Explanations>
            - Break down complex solutions into clear steps
            - Explain architectural decisions and tradeoffs
            - Provide performance analysis when relevant
            - Include debugging tips and common pitfalls
            - Add scaling considerations
            - Discuss alternative approaches
        </Explanations>

        <BestPractices>
            - Follow SOLID principles strictly
            - Implement appropriate design patterns
            - Consider security implications
            - Optimize for maintainability and readability
            - Follow 12-factor app principles
            - Implement CI/CD best practices
        </BestPractices>
    </ResponseBehavior>

    <Specialties>
        <LongCodeManagement>
            - Break down large codebases into manageable modules
            - Implement clear folder structures
            - Create comprehensive documentation
            - Use proper versioning and dependency management
            - Implement modular architecture
            - Create clear API boundaries
        </LongCodeManagement>

        <Debugging>
            - Provide sophisticated debugging strategies
            - Implement comprehensive logging
            - Add performance profiling
            - Include error tracking and reporting
            - Use debugging tools effectively
            - Implement monitoring solutions
        </Debugging>

        <Security>
            - Implement security best practices
            - Handle sensitive data properly
            - Prevent common vulnerabilities
            - Add input validation and sanitization
            - Implement authentication and authorization
            - Follow OWASP guidelines
        </Security>

        <DataManagement>
            - Design efficient database schemas
            - Optimize database queries
            - Implement caching strategies
            - Handle large-scale data processing
            - Implement data validation
            - Ensure data integrity
        </DataManagement>
    </Specialties>
</AIFramework>"""

MAX_TOKENS = 4096
TEMPERATURE = 0.7

class ChatState:
    """Manages the state of the chat session"""
    def __init__(self):
        self.client: Optional[anthropic.Client] = None
        self.system_prompt: str = DEFAULT_SYSTEM_PROMPT
        self.history: List[Tuple[str, str]] = []

    def initialize_client(self, api_key: str) -> Tuple[bool, str]:
        """Initialize the Anthropic client with the provided API key"""
        try:
            self.client = anthropic.Client(api_key=api_key)
            # Test the API key with a minimal request
            self.client.messages.create(
                model="claude-3-5-sonnet-20241022",
                max_tokens=10,
                messages=[{"role": "user", "content": "test"}]
            )
            return True, "API key validated successfully!"
        except Exception as e:
            self.client = None
            return False, f"Error validating API key: {str(e)}"

    def update_system_prompt(self, new_prompt: str) -> str:
        """Update the system prompt"""
        self.system_prompt = new_prompt
        return "System prompt updated successfully!"

    def generate_response(self, message: str) -> str:
        """Generate a response using the Claude API"""
        if not self.client:
            return "Please enter a valid API key first."
    
        try:
            # Convert history to message format
            messages = []
            for human_msg, assistant_msg in self.history:
                messages.extend([
                    {"role": "user", "content": human_msg},
                    {"role": "assistant", "content": assistant_msg}
                ])
            
            # Add the current message
            messages.append({"role": "user", "content": message})
    
            # Make API call with system parameter
            response = self.client.messages.create(
                model="claude-3-5-sonnet-20241022",
                max_tokens=MAX_TOKENS,
                temperature=TEMPERATURE,
                system=self.system_prompt,  # Pass system prompt as a separate parameter
                messages=messages
            )
    
            return response.content[0].text
    
        except Exception as e:
            return f"Error generating response: {str(e)}"

    def clear_history(self) -> None:
        """Clear the conversation history"""
        self.history = []

def create_demo() -> gr.Blocks:
    """Create the Gradio interface"""
    chat_state = ChatState()

    with gr.Blocks(title="Claude Coding Assistant", theme=gr.themes.Soft()) as demo:
        with gr.Row():
            # Main chat area (2/3 width)
            with gr.Column(scale=2):
                gr.Markdown("""
                # Claude Coding Assistant 🤖💻
                
                A specialized coding assistant powered by Claude 3.5 Sonnet. Features:
                - Complete code solutions with best practices
                - Detailed explanations and documentation
                - Code review and optimization suggestions
                - Security considerations and design patterns
                """)

                # API Key input section
                with gr.Row():
                    api_key_input = gr.Textbox(
                        label="Anthropic API Key",
                        placeholder="Enter your API key...",
                        type="password",
                        scale=4
                    )
                    validate_btn = gr.Button("Validate API Key", scale=1)
                    status_text = gr.Textbox(label="Status", interactive=False, scale=2)

                # Chat interface
                chatbot = gr.Chatbot(
                    height=600,
                    show_label=False,
                    container=True,
                    show_copy_button=True
                )
                
                with gr.Row():
                    message_input = gr.Textbox(
                        label="Your message",
                        placeholder="Ask me about coding, software design, or technical problems...",
                        lines=3,
                        scale=4
                    )
                    submit_btn = gr.Button("Send", scale=1)
                    clear_btn = gr.Button("Clear Chat", scale=1)

            # Sidebar for system prompt (1/3 width)
            with gr.Column(scale=1):
                gr.Markdown("### System Prompt Editor")
                system_prompt_input = gr.TextArea(
                    value=DEFAULT_SYSTEM_PROMPT,
                    label="Customize the AI's behavior",
                    lines=20,
                    placeholder="Enter system prompt..."
                )
                update_prompt_btn = gr.Button("Update System Prompt")
                prompt_status = gr.Textbox(label="Prompt Status", interactive=False)

        # Event handlers
        def validate_api_key(api_key: str) -> str:
            success, message = chat_state.initialize_client(api_key)
            return message

        def update_system_prompt(new_prompt: str) -> str:
            return chat_state.update_system_prompt(new_prompt)

        def respond(message: str, history: List[Tuple[str, str]]) -> Tuple[List[Tuple[str, str]], str]:
            if not chat_state.client:
                return history + [("", "Please validate your API key first.")], ""
            
            response = chat_state.generate_response(message)
            chat_state.history = history + [(message, response)]
            return chat_state.history, ""

        def clear_chat() -> Tuple[List[Tuple[str, str]], str]:
            chat_state.clear_history()
            return [], ""

        # Connect event handlers
        validate_btn.click(
            validate_api_key,
            inputs=[api_key_input],
            outputs=[status_text]
        )

        update_prompt_btn.click(
            update_system_prompt,
            inputs=[system_prompt_input],
            outputs=[prompt_status]
        )

        submit_btn.click(
            respond,
            inputs=[message_input, chatbot],
            outputs=[chatbot, message_input]
        )

        message_input.submit(
            respond,
            inputs=[message_input, chatbot],
            outputs=[chatbot, message_input]
        )

        clear_btn.click(
            clear_chat,
            outputs=[chatbot, message_input]
        )

        # Add examples
        gr.Examples(
            examples=[
                "Can you help me implement a binary search tree in Python with type hints?",
                "Review this code for security issues:\n```python\ndef process_user_input(data):\n    result = eval(data)\n    return result```",
                "Explain the SOLID principles with practical examples in TypeScript.",
                "How would you implement rate limiting in a REST API?",
            ],
            inputs=message_input,
            label="Example Questions"
        )

    return demo

# Create and launch the demo
demo = create_demo()

# For local testing
if __name__ == "__main__":
    demo.launch()