import streamlit as st
from bs4 import BeautifulSoup
from openai import OpenAI
from dotenv import load_dotenv
import os

# ==== Load API Key from .env ====
load_dotenv()
hf_api_key = os.getenv("deploy")

if not hf_api_key:
    st.error("❌ API key not found. Please set HUGGINGFACE_API_KEY in your .env file.")
    st.stop()

# ==== Initialize Client ====
client = OpenAI(
    base_url="https://router.huggingface.co/v1",
    api_key=hf_api_key,
)

# ==== Streamlit Page Setup ====
st.set_page_config(page_title="OSS ChatGPT", layout="wide")
st.title("🤖 ChatGPT")

# ==== Sidebar ====
st.sidebar.title("🛠️ Settings")
model_choice = st.sidebar.selectbox("Choose a model", [
    "openai/gpt-oss-20b:hyperbolic",
    "openai/gpt-oss-120b:hyperbolic"
])



if st.sidebar.button("🧹 Clear Chat"):
    st.session_state.messages = []
    st.rerun()

# ==== Session Initialization ====
if "messages" not in st.session_state:
    st.session_state.messages = []

# ==== LaTeX Rendering Helper ====
def render_markdown_with_latex(text: str):
    mathjax_script = """
    <script type="text/javascript">
      MathJax = {
        tex: {
          inlineMath: [['$', '$']],
          displayMath: [['$$', '$$']]
        },
        svg: {
          fontCache: 'global'
        }
      };
    </script>
    <script async src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-svg.js"></script>
    """
    st.markdown(mathjax_script, unsafe_allow_html=True)
    st.markdown(text, unsafe_allow_html=True)
# ... (keep all your existing imports and setup code until the message display section)
# ... (keep all your existing imports and setup code until the message display section)

# ==== Display Chat History ====
for i, msg in enumerate(st.session_state.messages):
    with st.chat_message(msg["role"]):
        if msg["role"] == "assistant":
            # Check if content contains thinking pattern
            if msg["content"].startswith("<think>") and "</think>" in msg["content"]:
                # Extract thinking and response parts
                thinking_part = msg["content"].split("<think>")[1].split("</think>")[0].strip()
                response_part = msg["content"].split("</think>")[1].strip()
                
                # Display the main response
                render_markdown_with_latex(response_part)
                
                # Create expander for thinking
                with st.expander("💭 Show Thinking"):
                    st.markdown(f"<span style='color: #666; font-style: italic'>{thinking_part}</span>", 
                               unsafe_allow_html=True)
            else:
                render_markdown_with_latex(msg["content"])
        else:
            # Display user messages
            render_markdown_with_latex(msg["content"])

# ==== Chat Input ====
st.write("Akshay")
if prompt := st.chat_input("Type your message...",key="unique_chat_input_key"):
    print(f"=> {prompt}")
    with st.chat_message("user"):
        st.markdown(prompt)
    
    # Add user message to history
    st.session_state.messages.append({"role": "user", "content": prompt})

    try:
        with st.chat_message("assistant"):
            response_placeholder = st.empty()

            with st.spinner("Thinking..."):
                completion = client.chat.completions.create(
                    model=model_choice,
                    messages=[
                        {"role": m["role"], "content": m["content"]}
                        for m in st.session_state.messages
                    ]
                )
                raw = completion.choices[0].message.content

                # Parse thinking and response
                if "<think>" in raw and "</think>" in raw:
                    thinking_part = raw.split("<think>")[1].split("</think>")[0].strip()
                    response_part = raw.split("</think>")[1].strip()

                    # Display main response
                    clean_response = BeautifulSoup(response_part, "html.parser").get_text()
                    with st.expander("💭 Show Thinking"):
                        st.markdown(f"<span style='color: #666; font-style: italic'>{thinking_part}</span>", 
                               unsafe_allow_html=True)
                    response_placeholder.markdown(clean_response)
                    # Store both parts in history
                    full_content = f"<think>{thinking_part}</think>{clean_response}"
                    
                else:
                    clean_response = BeautifulSoup(raw, "html.parser").get_text()
                    response_placeholder.markdown(clean_response)
                    full_content = clean_response

            # Add assistant response to history
            st.session_state.messages.append({"role": "assistant", "content": full_content})

    except Exception as e:
        error_msg = f"❌ Error: {str(e)}"
        with st.chat_message("assistant"):
            st.error(error_msg)
        st.session_state.messages.append({"role": "assistant", "content": error_msg})