# Import os to handle environment variables
#
# To run this app
# 
# streamlit run <path to this file, for example, rag\streamlit_app_basic.py>
#
import os
import uuid
import datetime
from dotenv import load_dotenv
import streamlit as st
from openai import AzureOpenAI

load_dotenv()

# --- Session State Initialization ---
if "conversations" not in st.session_state:
    st.session_state.conversations = {}

if "current_conversation_id" not in st.session_state:
    conversation_id = str(uuid.uuid4())
    st.session_state.conversations[conversation_id] = {
        "created_at": datetime.datetime.now().strftime("%Y-%m-%d %H:%M"),
        "messages": [
            {"role": "system", "content":
                """You are a helpful AI assistant which answers questions about cricket and sports in general.
                   Do not answer questions about other topics.
                   If you do not know the answer, say 'I do not know the answer to that question.'"""
            }
        ]
    }
    st.session_state.current_conversation_id = conversation_id

def get_current_conversation():
    return st.session_state.conversations[st.session_state.current_conversation_id]["messages"]

def generate_response(input_text):
    client = AzureOpenAI(
        azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
        api_key=os.getenv("AZURE_OPENAI_KEY"),
        api_version=os.getenv("AZURE_OPENAI_API_VERSION")
    )
    model_name = os.getenv("AZURE_OPENAI_MODEL_NAME")
    try:
        current_messages = get_current_conversation()
        current_messages.append({"role": "user", "content": input_text})
        with st.spinner("Generating Answer. Please wait..."):
            response = client.chat.completions.create(
                model=model_name,
                messages=current_messages
            )
        answer = response.choices[0].message.content.strip()
        current_messages.append({"role": "assistant", "content": answer})
        return answer
    except Exception as e:
        print(f"Error: {e}")
        st.error("An error occurred while processing your request. Please check your Azure configuration.")
        return None

def create_new_conversation():
    conversation_id = str(uuid.uuid4())
    st.session_state.conversations[conversation_id] = {
        "created_at": datetime.datetime.now().strftime("%Y-%m-%d %H:%M"),
        "messages": [
            {"role": "system", "content":
                """You are a helpful AI assistant which answers questions about cricket and sports in general.
                   Do not answer questions about other topics.
                   If you do not know the answer, say 'I do not know the answer to that question.'"""
            }
        ]
    }
    st.session_state.current_conversation_id = conversation_id
    st.session_state.user_question = ""  # Clear the input box

# --- Sidebar: Show only session titles ---
st.sidebar.title("Conversation Sessions")
for conv_id, conv_data in st.session_state.conversations.items():
    # Use first user message as title, else fallback to timestamp
    title = "New Conversation"
    for msg in conv_data["messages"]:
        if msg["role"] == "user":
            title = msg["content"][:30] + "..." if len(msg["content"]) > 30 else msg["content"]
            break
    if st.sidebar.button(title, key=f"conv_{conv_id}"):
        st.session_state.current_conversation_id = conv_id

# --- Main Page Layout ---
col1, col2 = st.columns([8, 1], gap="small")
with col2:
    st.markdown(
        """
        <style>
        div.stButton > button {
            width: 100%;
            background-color: #2563eb;
            color: white;
            font-weight: bold;
            border-radius: 6px;
            padding: 0.5em 0.8em;
            font-size: 1.1em;
        }
        </style>
        """,
        unsafe_allow_html=True,
    )
    st.button("🆕 New Chat", key="new_chat_btn", help="Start a new conversation", use_container_width=True, on_click=create_new_conversation)

st.markdown("<h1 style='text-align: center; color: #2563eb;'>Darshit's Chatbot 🤖</h1>", unsafe_allow_html=True)
st.markdown("---")

# --- Display Current Conversation ---
chat_container = st.container()
with chat_container:
    for message in get_current_conversation():
        if message["role"] == "user":
            st.markdown(
                f"<div style='background-color:#e0e7ff; border-radius:8px; padding:8px; margin-bottom:4px;'><b>You:</b> {message['content']}</div>",
                unsafe_allow_html=True,
            )
        elif message["role"] == "assistant":
            st.markdown(
                f"<div style='background-color:#f1f5f9; border-radius:8px; padding:8px; margin-bottom:8px;'><b>Bot:</b> {message['content']}</div>",
                unsafe_allow_html=True,
            )

# --- User Input ---
st.markdown("---")
user_question = st.text_input(
    "Ask any question about cricket or sports (Press Enter to submit):",
    key="user_question",
    placeholder="Type your question here...",
)
if user_question:
    answer = generate_response(user_question)
    if answer:
        st.info(answer)