import streamlit as st
from interpreter import interpreter
import os

# Page configuration
st.set_page_config(page_title="AutoInterpreter", layout="wide")


# Initialize session state for settings if not exists
if "settings" not in st.session_state:
    st.session_state.settings = {
        "api_key": os.getenv("HF_API_KEY", ""),
        "api_base": "https://api-inference.huggingface.co/models/Qwen/Qwen2.5-72B-Instruct",
        "model": "huggingface/Qwen/Qwen2.5-72B-Instruct",
        "auto_run": True,
        "context_window": 8000,
        "max_tokens": 4000
    }

# Create header with title and settings button
col1, col2 = st.columns([0.9, 0.1])
with col1:
    st.markdown("# Autointerpreter")
    st.markdown("Run Any Code. The Final AI Coding Experience.")
with col2:
    settings_button = st.button("⚙️", help="Settings")

# Settings modal
if settings_button:
    settings_modal = st.container()
    with settings_modal:
        st.markdown("### Settings")
        cols = st.columns(2)
        
        with cols[0]:
            # API Settings
            st.text_input(
                "API Key", 
                value=st.session_state.settings["api_key"],
                type="password",
                key="api_key",
                on_change=lambda: st.session_state.settings.update({"api_key": st.session_state.api_key})
            )
            st.text_input(
                "Model", 
                value=st.session_state.settings["model"],
                key="model",
                on_change=lambda: st.session_state.settings.update({"model": st.session_state.model})
            )
        
        with cols[1]:
            # Model Settings
            st.toggle(
                "Auto Run", 
                value=st.session_state.settings["auto_run"],
                key="auto_run",
                on_change=lambda: st.session_state.settings.update({"auto_run": st.session_state.auto_run})
            )
            st.number_input(
                "Max Tokens", 
                value=st.session_state.settings["max_tokens"],
                min_value=100,
                max_value=8000,
                key="max_tokens",
                on_change=lambda: st.session_state.settings.update({"max_tokens": st.session_state.max_tokens})
            )

# Apply settings to interpreter
interpreter.llm.api_key = st.session_state.settings["api_key"]
interpreter.llm.api_base = st.session_state.settings["api_base"]
interpreter.llm.model = st.session_state.settings["model"]
interpreter.auto_run = st.session_state.settings["auto_run"]
interpreter.context_window = st.session_state.settings["context_window"]
interpreter.max_tokens = st.session_state.settings["max_tokens"]

# Initialize messages session state
if "messages" not in st.session_state:
    st.session_state.messages = []

# Clear button
if st.button("🗑️ Clear", help="Clear chat"):
    interpreter.messages = []
    st.session_state.messages = []
    st.rerun()

# Display chat history
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# User input
user_input = st.chat_input("Enter your message:")

if user_input:
    # Display user message
    st.chat_message("user").write(user_input)
    st.session_state.messages.append({"role": "user", "content": user_input})
    
    try:
        # Create a chat message container for the assistant
        with st.chat_message("assistant"):
            response_placeholder = st.empty()
            message_buffer = []
            code_buffer = []
            
            # Stream the response
            for chunk in interpreter.chat(user_input, stream=True):
                if isinstance(chunk, dict):
                    content = chunk.get('content')
                    if content is not None and not any(skip in str(content) for skip in ["context window", "max_tokens", "<|im_end|>"]):
                        if chunk.get('type') == 'console':
                            # Accumulate code separately
                            code_buffer.append(str(content))
                            # Show complete message + current code
                            full_response = []
                            if message_buffer:
                                full_response.extend(message_buffer)
                            if code_buffer:
                                full_response.append(f"\n```python\n{''.join(code_buffer)}\n```\n")
                            response_placeholder.markdown(''.join(full_response))
                        else:
                            # Accumulate message until we have a complete thought
                            current = str(content)
                            message_buffer.append(current)
                            if '.' in current or '\n' in current or len(''.join(message_buffer)) > 80:
                                # Show complete message + current code
                                full_response = []
                                if message_buffer:
                                    full_response.extend(message_buffer)
                                if code_buffer:
                                    full_response.append(f"\n```python\n{''.join(code_buffer)}\n```\n")
                                response_placeholder.markdown(''.join(full_response))
            
            # Store the complete response
            final_response = []
            if message_buffer:
                final_response.extend(message_buffer)
            if code_buffer:
                final_response.append(f"\n```python\n{''.join(code_buffer)}\n```\n")
            
            st.session_state.messages.append({
                "role": "assistant", 
                "content": ''.join(final_response)
            })
            
    except Exception as e:
        st.error(f"Error: {str(e)}")