import streamlit as st
import os
import time
from services.llm import initialize_llm, initialize_embeddings
from services.vector_store import create_vector_store, retrive_vector_store, generate_prompt
from services.pdf_processing import load_and_split_pdf
from utils.helpers import extract_thoughts, response_generator
import subprocess

# try:
#     print("🚀 Checking and starting Ollama...")
#     subprocess.run(["bash", "install_ollama.sh"], check=True)
#     print("✅ Ollama is running!")
# except subprocess.CalledProcessError as e:
#     print(f"❌ Error: {e}")


# Custom CSS for chat styling
CHAT_CSS = """
    <style>
        .user-message {
            text-align: right;
            background-color: #3c8ce7;
            color: white;
            padding: 10px;
            border-radius: 10px;
            margin-bottom: 10px;
            display: inline-block;
            width: fit-content;
            max-width: 70%;
            margin-left: auto;
            box-shadow: 0px 4px 6px rgba(0, 0, 0, 0.1);
        }
        .assistant-message {
            text-align: left;
            background-color: #d16ba5;
            color: white;
            padding: 10px;
            border-radius: 10px;
            margin-bottom: 10px;
            display: inline-block;
            width: fit-content;
            max-width: 70%;
            margin-right: auto;
            box-shadow: 0px 4px 6px rgba(0, 0, 0, 0.1);
        }
    </style>
"""

# Streamlit UI Setup
st.set_page_config(page_title="DocChatAI", layout="wide")
st.title("📄 DocChatAI | Chat Using Documents")
# Expandable Disclaimer Section
with st.expander("⚠️ Disclaimer (Click to expand)"):
    st.markdown("""
    - This AI chatbot provides **informational responses only** and should not be used as **legal, medical, or financial advice**.
    - The accuracy of responses depends on the provided **context and training data**.
    - **Use at your own discretion** and always verify important information from reliable sources.
    """)

# Sidebar
st.sidebar.title("DocChatAI")
st.sidebar.subheader("Chat using PDF Document")
st.sidebar.write("---")

# Model Selection
selected_model = st.sidebar.radio("Choose Model", ["deepseek-r1:1.5b"])
st.sidebar.write("---")

# Hyperparameters
temperature = st.sidebar.slider("Temperature", 0.0, 1.0, 0.7, 0.1)
top_p = st.sidebar.slider("Top-p (Nucleus Sampling)", 0.0, 1.0, 0.9, 0.05)
max_tokens = st.sidebar.number_input("Max Tokens", 10, 2048, 1024, 10)
st.sidebar.write("---")

# File Upload
uploaded_file = st.sidebar.file_uploader("📂 Upload a PDF", type=["pdf"])
st.sidebar.write("---")

# About Section
st.sidebar.write("📌 **About Me**")
st.sidebar.write("👤 **Name:** Deepak Yadav")
st.sidebar.write("💡 **Bio:** Passionate about AI and Machine Learning.")
st.sidebar.markdown("[GitHub](https://github.com/deepak7376) | [LinkedIn](https://www.linkedin.com/in/dky7376/)")
st.sidebar.write("---")

# Initialize LLM
llm = initialize_llm(selected_model, temperature, top_p, max_tokens)
embeddings = initialize_embeddings()

# Document Handling
retriever = None
if uploaded_file:
    os.makedirs("docs", exist_ok=True)
    filepath = os.path.join("docs", uploaded_file.name)

    with open(filepath, "wb") as f:
        f.write(uploaded_file.read())

    # Load and process PDF
    splits = load_and_split_pdf(filepath)
    vectorstore = create_vector_store(splits, embeddings)
    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})

# Apply custom CSS
st.markdown(CHAT_CSS, unsafe_allow_html=True)

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

# Display previous messages
for message in st.session_state.messages:
    if message['thinking_part']:
        with st.expander("💭 Thought Process"):
                st.markdown(message['thinking_part'])
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Chat Input
if user_input := st.chat_input("💬 Ask something..."):
    st.session_state.messages.append({"role": "user", "content": user_input, "thinking_part": False})

    with st.chat_message("user"):
        st.markdown(user_input)

    # Measure response time
    start_time = time.time()

    # Generate response
    context = retrive_vector_store(retriever, user_input) if retriever else "No context"
    query = generate_prompt(context=context, question=user_input)
    # response = llm.invoke(query)

    response = llm.create_chat_completion(
	messages = [
		{
			"role": "user",
			"content": f"{query}"
		}
            ]
        )

    # Calculate response time
    response_time = round(time.time() - start_time, 2)

    # Extract thoughts and main answer
    thinking_part, main_answer = extract_thoughts(response['choices'][0]['message']['content'])

    # Display AI response
    with st.chat_message("assistant"):
        if thinking_part:
            with st.expander("💭 Thought Process"):
                st.markdown(thinking_part)

        # **Formatted Response Display**
        formatted_response = f""" 
        {main_answer}  

        ⏳ **Response Time:** {response_time} seconds  
        """
        st.markdown(formatted_response, unsafe_allow_html=True)

    # Save to session history
    st.session_state.messages.append({"role": "assistant", "content": formatted_response, "thinking_part": thinking_part})