# ----- ----- IMPORTS
import streamlit as st
import requests, time, regex, re
from datetime import datetime
from dotenv import load_dotenv
import certifi
import os

load_dotenv()

from ibm_watsonx_ai import Credentials, APIClient
from ibm_watsonx_ai.foundation_models import ModelInference
from jinja2 import Template

from src.bot_specs import bot_name, bot_icon, user_icon
from src.parameters import (
    model_id,
    system_prompt,
    params,
    display_chat_history,
    stream_outputs,
    app_password,
    wx_api_key,
    wx_project_id,
    wx_url,
    info_tag
)
from src.helper_functions import (
    setup_watsonxai_client,
    create_pdf_from_chat,
    watsonx_chat_prompt,
    generate_response,
)

# ----- ----- PAGE CONFIG
st.set_page_config(
    page_title=bot_name,
    page_icon=bot_icon,
    initial_sidebar_state="collapsed",
    layout="centered",
)


def check_password():
    def password_entered():
        if st.session_state["password"] == app_password:
            st.session_state["password_correct"] = True
            del st.session_state["password"]
        else:
            st.session_state["password_correct"] = False

    if "password_correct" not in st.session_state:
        st.markdown("\n\n")
        st.text_input(
            "Enter the password",
            type="password",
            on_change=password_entered,
            key="password",
        )
        st.divider()
        st.info(info_tag)
        return False
    elif not st.session_state["password_correct"]:
        st.markdown("\n\n")
        st.text_input(
            "Enter the password",
            type="password",
            on_change=password_entered,
            key="password",
        )
        st.divider()
        st.info(info_tag)
        st.error("😕 Password incorrect")
        return False
    else:
        return True


if not check_password():
    st.stop()

if "current_page" not in st.session_state:
    st.session_state.current_page = 0


def initialize_session_state():
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []

wx_client = setup_watsonxai_client(
    api_key=wx_api_key, project_id=wx_project_id, url=wx_url
)


# ----- ----- MAIN APP
def main():
    initialize_session_state()
    st.subheader(f"{bot_name} {bot_icon}")

    if display_chat_history:
        for message in st.session_state.chat_history:
            with st.chat_message(
                message["role"],
                avatar=user_icon if message["role"] == "user" else bot_icon,
            ):
                st.markdown(message["content"])
    user_input = st.chat_input("You:", key="user_input")

    if user_input:
        # Add user message to chat history
        st.session_state.chat_history.append({"role": "user", "content": user_input})
        with st.chat_message("user", avatar=user_icon):
            st.markdown(user_input)

        with st.chat_message(bot_name, avatar=bot_icon):
            # Build messages with baseline + chat history
            messages = [{"role": "system", "content": system_prompt}]

            messages.extend(st.session_state.chat_history)

            stream_generator = watsonx_chat_prompt(
                messages=messages,
                stream=stream_outputs,
                client=wx_client,
                params=params,
                model_id=model_id,
            )

            text_output = generate_response(stream_generator, stream=stream_outputs)

            # Stream the response with typewriter effect
            assistant_response = st.write_stream(text_output)

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

    if st.session_state.chat_history:
        now = datetime.now()
        date_str = now.strftime("%Y-%m-%d")
        try:
            pdf_buffer = create_pdf_from_chat(st.session_state.chat_history)
            st.download_button(
                label="Download Chat History as PDF",
                data=pdf_buffer,
                file_name=f"chat_history_{date_str}.pdf",
                mime="application/pdf",
            )
        except Exception as e:
            st.error(f"An error occurred while generating the PDF: {str(e)}")


if __name__ == "__main__":
    main()