import os
import re
import streamlit as st
from openai import OpenAI
from dotenv import load_dotenv
from docx import Document
from io import BytesIO

# Load environment variables
load_dotenv()

# Page config
st.set_page_config(page_title="🎓 Sharqi Education Chatbot", layout="wide")
st.title("🎓 Sharqi Education Chatbot")

# Check API key
groq_api_key = os.getenv("GROQ_API_KEY")
if not groq_api_key:
    st.error("❌ GROQ_API_KEY is not set in the environment variables.")
    st.stop()

# OpenAI client for GROQ
client = OpenAI(
    api_key=groq_api_key,
    base_url="https://api.groq.com/openai/v1"
)

model_name = "llama3-70b-8192"

# Session state
if "history" not in st.session_state:
    st.session_state.history = []

# Function to build chat messages
def build_messages():
    return [
        {"role": "user" if sender == "You" else "assistant", "content": msg}
        for sender, msg in st.session_state.history
    ]

# Clean model output
def clean_think_tags(text: str) -> str:
    return re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL).strip()

# Detect Urdu (rudimentary check)
def is_urdu(text: str) -> bool:
    return any("\u0600" <= char <= "\u06FF" for char in text)

# Convert text to Word and return BytesIO
def create_word_file(text):
    doc = Document()
    if is_urdu(text):
        style = doc.styles['Normal']
        font = style.font
        font.name = 'Jameel Noori Nastaleeq'
    doc.add_paragraph(text)
    buffer = BytesIO()
    doc.save(buffer)
    buffer.seek(0)
    return buffer

# Show previous chat history
for sender, message in st.session_state.history:
    with st.chat_message("user" if sender == "You" else "assistant"):
        align = "right" if is_urdu(message) else "left"
        st.markdown(f"<div style='text-align:{align}; direction:{'rtl' if align=='right' else 'ltr'};'>{message}</div>", unsafe_allow_html=True)

# Chat input
user_input = st.chat_input("اردو یا انگلش میں سوال کریں...")

if user_input:
    st.session_state.history.append(("You", user_input))
    with st.chat_message("user"):
        st.markdown(f"<div style='text-align:right;'>{user_input}</div>", unsafe_allow_html=True)

    with st.chat_message("assistant"):
        placeholder = st.empty()
        placeholder.write("⏳ Thinking...")

        try:
            response = client.chat.completions.create(
                model=model_name,
                messages=build_messages()
            )
            raw_output = response.choices[0].message.content
            reply = clean_think_tags(raw_output)
        except Exception as e:
            reply = f"❌ API Error: {e}"

        placeholder.markdown(
            f"<div style='text-align:{'right' if is_urdu(reply) else 'left'};'>{reply}</div>", unsafe_allow_html=True
        )
        st.session_state.history.append(("Bot", reply))

        # Buttons for copy & download
        with st.expander("🔧 Options"):
            st.code(reply, language="text")
            st.download_button(
                label="📄 Download Reply as Word",
                data=create_word_file(reply),
                file_name="reply.docx",
                mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document"
            )
            st.text_input("📋 Copy this:", value=reply, key="copy_field")