import streamlit as st
import time
from backend import BDAgent
import boto3
from boto3.dynamodb.conditions import Key

# Configure the page - MUST BE FIRST
st.set_page_config(
    page_title="Chatbot",
    page_icon="💬",
    layout="centered"
)

session = boto3.Session()
dynamodb = boto3.resource("dynamodb", region_name="us-east-1")

def list_bases():
    table = dynamodb.Table("poc_dnx_monthly_summary")
    response = table.query(
        IndexName="item_type_index",
        KeyConditionExpression=Key("item_type").eq("contexto")
    )
    return [item["id"].removesuffix("_contexto") for item in response.get("Items", [])]

# Model selection with session state persistence
MODELS = [
    "anthropic.claude-haiku-4-5-20251001-v1:0",
    "anthropic.claude-sonnet-4-5-20250929-v1:0",
    "meta.llama4-maverick-17b-instruct-v1:0",
    "meta.llama4-scout-17b-instruct-v1:0",
    "amazon.nova-lite-v1:0",
    "amazon.nova-pro-v1:0",
    "amazon.nova-2-lite-v1:0"
]

selected_value = st.selectbox(
    "Selecione o modelo:",
    MODELS,
    key="selected_model"
)
BASES = list_bases()
BASES.sort()
base = st.selectbox(
    "Selecione a base:",
    BASES,
    key="selected_base"
)


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

# Display title
st.title("💬 Chatbot")

# Display chat history
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 history
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # Display user message
    with st.chat_message("user"):
        st.markdown(prompt)
    
    # Display assistant response
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        
        # Simulate streaming response (replace with actual API call)

        result = BDAgent.main(prompt,str(st.session_state.messages),selected_value,base)
        full_response = result["response"]

        # Simulate typing effect
        displayed_response = ""
        for char in full_response:
            displayed_response += char
            message_placeholder.markdown(displayed_response + "▌")
            time.sleep(0.01)

        message_placeholder.markdown(full_response)
        st.caption(f"Tokens: {result['input_tokens']} in / {result['output_tokens']} out / {result['total_tokens']} total")

    # Add assistant response to chat history
    st.session_state.messages.append({"role": "assistant", "content": full_response})

# Add a sidebar with options
with st.sidebar:
    st.header("Options")
    
    if st.button("Clear Chat History"):
        st.session_state.messages = []
        st.rerun()
    
    st.divider()
    
    st.markdown("""
    ### How to use:
    1. Type your message in the input box
    2. Press Enter or click Send
    3. View the conversation history
    """)