import streamlit as st
import anthropic
import json
import os
from utils import ANTHROPIC_MODEL_OPTIONS

def display():
    """
    Display the Anthropic models tab.
    """
    st.header("Anthropic (Claude) Models")

    # API key input (with warning about security)
    anthropic_key = st.text_input(
        "Enter your Anthropic API Key",
        type="password",
        help="⚠️ Never share your API key. Leave empty to use ANTHROPIC_API_KEY environment variable.",
    )

    # If no key provided, try to get from environment
    if not anthropic_key:
        anthropic_key = os.environ.get("ANTHROPIC_API_KEY", "")

    # Model selection for Anthropic
    selected_anthropic_model = st.selectbox(
        "Select Claude Model", list(ANTHROPIC_MODEL_OPTIONS.keys())
    )

    # System message (optional)
    st.subheader("System Message (Optional)")
    system_message = st.text_area(
        "System Message", placeholder="e.g., You are a helpful assistant", height=100
    )

    # User message input
    st.subheader("Message Content")
    anthropic_user_message = st.text_area(
        "Enter your message here",
        placeholder="Hello, Claude! How are you today?",
        height=200,
        key="anthropic_message",
    )

    # Button to count tokens for Anthropic
    if st.button("Count Tokens (Anthropic)"):
        if not anthropic_key:
            st.error(
                "No Anthropic API key found. Please enter a key or set the ANTHROPIC_API_KEY environment variable."
            )
        elif not anthropic_user_message:
            st.warning("Please enter a message to count tokens")
        else:
            try:
                # Initialize client with API key
                client = anthropic.Anthropic(api_key=anthropic_key)

                # Create the request
                count_request = {
                    "model": ANTHROPIC_MODEL_OPTIONS[selected_anthropic_model],
                    "messages": [{"role": "user", "content": anthropic_user_message}],
                }

                # Add system message if provided
                if system_message:
                    count_request["system"] = system_message

                # Make the API call to count tokens
                response = client.messages.count_tokens(**count_request)

                # Display results
                st.success(f"Input tokens: {response.input_tokens}")

                # Display the full JSON response in an expandable section
                with st.expander("View Full API Response"):
                    st.code(
                        json.dumps(response.model_dump(), indent=2), language="json"
                    )

            except Exception as e:
                st.error(f"An error occurred: {str(e)}")