"""
Simple Chatbot using Streamlit and Anthropic Claude
==================================================

A clean, interactive chatbot interface powered by Claude AI.
"""

import streamlit as st
import anthropic
from typing import List, Dict
import os

# Configure the page
st.set_page_config(
    page_title="Claude Chatbot",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

def initialize_session_state():
    """Initialize session state variables."""
    if "messages" not in st.session_state:
        st.session_state.messages = []
    if "client" not in st.session_state:
        st.session_state.client = None

def setup_anthropic_client(api_key: str) -> anthropic.Anthropic:
    """Setup and return Anthropic client."""
    try:
        client = anthropic.Anthropic(api_key=api_key)
        return client
    except Exception as e:
        st.error(f"Error setting up Anthropic client: {str(e)}")
        return None

def get_claude_response(client: anthropic.Anthropic, messages: List[Dict], model: str = "claude-3-haiku-20240307") -> str:
    """Get response from Claude API."""
    try:
        # Convert Streamlit messages to Anthropic format
        anthropic_messages = []
        for msg in messages:
            if msg["role"] in ["user", "assistant"]:
                anthropic_messages.append({
                    "role": msg["role"],
                    "content": msg["content"]
                })
        
        response = client.messages.create(
            model=model,
            max_tokens=1000,
            temperature=0.7,
            messages=anthropic_messages
        )
        
        return response.content[0].text
    
    except Exception as e:
        return f"Error getting response: {str(e)}"

def main():
    """Main application function."""
    initialize_session_state()
    
    # Sidebar for configuration
    with st.sidebar:
        st.title("🤖 Claude Chatbot")
        st.markdown("---")
        
        # API Key input
        api_key = st.text_input(
            "Anthropic API Key",
            type="password",
            help="Enter your Anthropic API key to start chatting",
            placeholder="sk-ant-..."
        )
        
        # Model selection
        model = st.selectbox(
            "Select Claude Model",
            options=[
                "claude-3-haiku-20240307",
                "claude-3-sonnet-20240229",
                "claude-3-opus-20240229"
            ],
            index=0,
            help="Choose the Claude model to use"
        )
        
        # Clear chat button
        if st.button("🗑️ Clear Chat", use_container_width=True):
            st.session_state.messages = []
            st.rerun()
        
        st.markdown("---")
        st.markdown("""
        ### About
        This chatbot uses Anthropic's Claude AI to provide intelligent responses.
        
        **Features:**
        - Multiple Claude models
        - Persistent conversation history
        - Clean, responsive interface
        
        **Usage:**
        1. Enter your Anthropic API key
        2. Start typing your message
        3. Enjoy chatting with Claude!
        """)
    
    # Main chat interface
    st.title("💬 Chat with Claude")
    
    # Check if API key is provided
    if not api_key:
        st.info("👈 Please enter your Anthropic API key in the sidebar to start chatting.")
        st.markdown("""
        ### Getting Started
        
        1. **Get an API Key**: Visit [Anthropic Console](https://console.anthropic.com/) to get your API key
        2. **Enter the Key**: Paste it in the sidebar input field
        3. **Start Chatting**: Type your message and press Enter!
        
        ### Example Questions to Try:
        - "Explain quantum computing in simple terms"
        - "Write a Python function to calculate fibonacci numbers"
        - "What are the benefits of renewable energy?"
        - "Help me plan a weekend trip to San Francisco"
        """)
        return
    
    # Setup client
    if st.session_state.client is None:
        with st.spinner("Setting up connection to Claude..."):
            st.session_state.client = setup_anthropic_client(api_key)
    
    if st.session_state.client is None:
        st.error("Failed to setup Anthropic client. Please check your API key.")
        return
    
    # Display chat messages
    chat_container = st.container()
    with chat_container:
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])
    
    # Chat input
    if prompt := st.chat_input("Type your message here..."):
        # Add user message to chat
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        # Display user message immediately
        with chat_container:
            with st.chat_message("user"):
                st.markdown(prompt)
        
        # Get and display assistant response
        with chat_container:
            with st.chat_message("assistant"):
                message_placeholder = st.empty()
                with st.spinner("Claude is thinking..."):
                    response = get_claude_response(
                        st.session_state.client, 
                        st.session_state.messages,
                        model
                    )
                message_placeholder.markdown(response)
        
        # Add assistant response to chat history
        st.session_state.messages.append({"role": "assistant", "content": response})
    
    # Display conversation stats
    if st.session_state.messages:
        st.sidebar.markdown("---")
        st.sidebar.markdown(f"**Messages:** {len(st.session_state.messages)}")
        st.sidebar.markdown(f"**Model:** {model}")

if __name__ == "__main__":
    main()