import os
import datetime

import google.generativeai as genai
import streamlit as st 

from dotenv import load_dotenv
from typing import Any, Dict, List 

from db_ops import get_recent_conversations, init_db, store_conversation_summary, store_message 

# Loading environment variables from .env file 
load_dotenv()

# Configure Gemini API 
gemini_api_key = os.getenv("GEMINI_API_KEY")
genai.configure(api_key=gemini_api_key)

# Intialize the gemini Model 
model = genai.GenerativeModel("gemini-2.0-flash")


def chat_with_gemini(prompt:str, chat_history:List[Dict[str, str]]) -> str:
    try:
        # Format the message for Gemini 
        messages = []
        for msg in chat_history:
            if msg["role"] == "user":
                messages.append({"role": "user", "parts": msg["content"]})
            else:
                messages.append({"role": "model", "parts": [msg["content"]]})

        # Add current prompt 
        messages.append({"role": "user", "parts": [prompt]}) 

        # Generate the response from Gemini 
        chat = model.start_chat(history=messages[:-1])
        response = chat.send_message(prompt)
        return response.text 
    except Exception as e:
        return f"Error communicating with GEMINI API: {str(e)}" 
    

def summarize_conversations(conversations:List[Dict[str, Any]])->str:
    if not conversations:
        return "No recent conversations to summarize."
    
    # Format conversations for the model
    conversation_texts = []
    for idx, conv in enumerate(conversations, 1):
        conv_text = f"Conversation {idx} ({conv['timestamp']}):\n"
        for msg in conv["messages"]:
            conv_text += f"{msg['role'].upper()}: {msg['content']}\n"
        conversation_texts.append(conv_text)

    prompt = f"""
    Please provide a concise summary of the following recent conversations:

    {"\n\n".join(conversation_texts)}

    Focus on key topics discussed, questions asked, and information provided.
    Highlight any recurring themes or import points.
    """

    response = model.generate_content(prompt.strip())
    return response.text


# Streamlit UI 
def main():
    st.set_page_config(page_title="Gemini-Chatbot", page_icon="🤖")
    st.title("🤖 Gemini AI Chatbot")

    # Intialize the database 
    init_db()

    # Intialize session state for chat history and session ID
    if "session_id" not in st.session_state:
        st.session_state.session_id = datetime.datetime.now().strftime("%Y%m%d%H%M%S")

    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []

    # Chat input area 
    with st.container():
        user_input = st.chat_input("Type your message here...") 

        if user_input:
            # Add user message to chat history
            st.session_state.chat_history.append({"role": "user", "content": user_input})
            store_message(st.session_state.session_id, "user", user_input)

            # Get response from Gemini AI
            with st.spinner("Thinking..."):
                response = chat_with_gemini(user_input, st.session_state.chat_history)

            # Add assistant message to chat history
            st.session_state.chat_history.append({"role": "assistant", "content": response})
            store_message(st.session_state.session_id, "assistant", response)

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

    # Sidebar for recent conversations
    with st.sidebar:
        st.title("Conversation Recall")

        if st.button("🌸 Summarize Recent conversations"):
            with st.spinner("Generating summary..."):
                # Get recent conversations 
                recent_convs = get_recent_conversations(st.session_state.session_id, limit=5)

                # Generate summary 
                summary = summarize_conversations(recent_convs)

                # Store summary for recent conversations 
                if recent_convs:
                    store_conversation_summary(st.session_state.session_id, recent_convs[0]["id"], summary)

                # Display summary
                st.subheader("Summary of Recent Conversations")
                st.write(summary)

        # Clear chat button 
        if st.button("🗑️ Clear Chat"):
            st.session_state.chat_history = []
            st.session_state.session_id = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
            st.success("Chat history cleared!")
            st.rerun()

if __name__ == "__main__":
    main()