import os
import base64
import streamlit as st
from langchain_core.messages import HumanMessage, AIMessage

def get_base64_image(image_path):
    with open(image_path, 'rb') as img_file:
        return base64.b64encode(img_file.read()).decode()

def create_interface(call_model):
    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []

    # === Background Image ===
    background_image_path = os.path.join(os.getcwd(), "src", "interface", "assets", "uwf-page-edited.jpg")
    print("DEBUG: Looking for background image at:", background_image_path)
    print("DEBUG: File exists?", os.path.exists(background_image_path))

    if os.path.exists(background_image_path):
        base64_background_image = get_base64_image(background_image_path)
        st.markdown(
            f"""
            <style>
            .stApp {{
                background-image: url("data:image/jpeg;base64,{base64_background_image}");
                background-size: cover;
                background-position: center;
                background-repeat: no-repeat;
                background-attachment: fixed;
            }}
            h1 {{ color: #000000; }}
            p {{ color: #000000; }}
            .custom-caption {{ color: #A5A4A4; font-size: 16px; text-align: center; }}
            .chat-bubble {{
                max-width: 70%;
                padding: 10px;
                margin: 5px 0;
                border-radius: 10px;
                color: #000000;
            }}
            .human-message {{ background-color: #8DC8E8; align-self: flex-end; text-align: right; }}
            .ai-message {{ background-color: #E7E7E7; align-self: flex-start; }}
            </style>
            """,
            unsafe_allow_html=True
        )

    # === Header Section ===
    col1, col2 = st.columns([1, 3])
    with col1:
        image_path = os.path.join(os.getcwd(), "src", "interface", "assets", "argie-argonaut-gif-edited.png")
        print("DEBUG: Looking for logo image at:", image_path)
        print("DEBUG: File exists?", os.path.exists(image_path))
        if os.path.exists(image_path):
            st.image(image_path, width=200)
    with col2:
        st.markdown("<h1 style='color:black;'>ARGObot: UWF's Custom Q&A Chatbot</h1>", unsafe_allow_html=True)

    # === Example Prompts ===
    st.markdown("""
        <div style='text-align: center; font-size: 20px; font-weight: normal;'>
            <p style='font-size: 20px;'>Ask ARGObot a variety of questions based on the Student Handbook.</p>
        </div>
    """, unsafe_allow_html=True)

    st.write("Example topics:")
    st.markdown("<p style='color: #000000;'>- UWF's values</p>", unsafe_allow_html=True)
    st.markdown("<p style='color: #000000;'>- Student Rights and Responsibilities</p>", unsafe_allow_html=True)
    st.markdown("<p style='color: #000000;'>- UWF Policies and Regulations</p>", unsafe_allow_html=True)
    st.markdown("<p style='color: #000000;'>- UWF Appeals and Student Grievance Processes</p>", unsafe_allow_html=True)
    st.markdown("<p style='color: #000000;'>- Student Health and Wellbeing</p>", unsafe_allow_html=True)
    st.markdown("<p style='color: #000000;'>- Student Resources</p>", unsafe_allow_html=True)

    # === Chat Interface ===
    chat_container = st.container()
    with chat_container:
        st.markdown("<div class='chat-container'>", unsafe_allow_html=True)
        for message in st.session_state.chat_history:
            if isinstance(message, HumanMessage):
                st.markdown(f"<div class='chat-bubble human-message'><strong>You:</strong> {message.content}</div>", unsafe_allow_html=True)
            elif isinstance(message, AIMessage):
                st.markdown(f"<div class='chat-bubble ai-message'><strong>Chatbot:</strong> {message.content}</div>", unsafe_allow_html=True)
        st.markdown("</div>", unsafe_allow_html=True)

    # === Input Form ===
    with st.form(key='question_form', clear_on_submit=True):
        user_input = st.text_input("Enter your question here:", key="user_input")
        submit_button = st.form_submit_button(label='Submit')

        if submit_button and user_input:
            state = {"input": user_input, "chat_history": st.session_state.chat_history, "context": "", "answer": ""}
            result = call_model(state)

            if not result.get("answer"):
                result["answer"] = "Sorry, I couldn't generate an answer."

            st.session_state.chat_history.append(HumanMessage(user_input))
            st.session_state.chat_history.append(AIMessage(result["answer"]))
            st.rerun()
        elif submit_button:
            st.warning("Please enter a question.")