import os
import re
import gradio as gr
from langchain.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from langchain_chroma import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings

# Load embedding model and vector store from persisted DB
embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vector_store = Chroma(
    embedding_function=embedding_model,
    persist_directory="geometry_db",  # relative folder inside your Hugging Face Space
    collection_name="geometry_sol"
)

# Load OpenAI key (set this in Hugging Face Space Secrets)
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")

# Load the LLM (GPT-3.5)
llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.3)

# Unified prompt to auto-detect intent
template = PromptTemplate(
    input_variables=["context", "query"],
    template="""
You are a Virginia high school Geometry assistant. Based on the user question below, determine the correct response type and answer accordingly:

User Question: 
{query}

Based on the following SOL text:
{context}

Response Rules:
- If the question is asking for an **SOL number**, respond with:
  1. The exact SOL code (e.g., G.RLT.1)
  2. The exact description line from the SOL guide  
  ⚠️ Do not summarize. Only copy directly from the context.

- If the user asks for a **lesson plan**, provide:
  - Simple explanation of the concept
  - Real-world example
  - Engaging class activity  
  Format the output clearly with bullet points.

- If the user asks for a **worksheet**, include:
  - Concept summary
  - A worked example
  - 3 practice problems  
  Format the output clearly with bullet points.

- If the user asks for **proofs**, include:
  - Student-friendly explanation
  - Real-world connection
  - One short class activity  
  Format the output clearly with bullet points.

- If the user asks for **flashcards**, generate 5 cards, each with:
  - A clear question
  - A short answer  
  Format the output clearly with bullet points.

Only answer one way depending on the intent of the question.
"""
)

# Optional: shortcut to solve simple math problems (like area of rectangle)
def try_math_solver(query):
    match = re.search(r"rectangle.*l\s*=\s*(\d+).+w\s*=\s*(\d+)", query.lower())
    if match:
        l, w = int(match.group(1)), int(match.group(2))
        return f"The area of the rectangle is {l} × {w} = {l * w} square units."
    return None

# RAG function using unified intent-aware prompt
def rag_query(query):
    docs = vector_store.similarity_search(query, k=2)
    context = "\n\n".join([doc.page_content for doc in docs])
    prompt = template.format_prompt(context=context, query=query).to_string()
    return llm.invoke(prompt).content

# Gradio app function
def ask_geometry_sol(query):
    math_result = try_math_solver(query)
    if math_result:
        return math_result
    try:
        return rag_query(query)
    except Exception as e:
        return f"⚠️ Error: {type(e).__name__} - {str(e)}"

# Gradio UI (no need for manual response type selection anymore!)
iface = gr.Interface(
    fn=ask_geometry_sol,
    inputs=gr.Textbox(label="Enter your Geometry SOL question or topic"),
    outputs="text",
    title="📘 Virginia Geometry SOL Assistant",
    description="Ask about any 2023 Geometry SOL (Standards of Learning). The assistant will auto-detect if you want a lesson plan, worksheet, proof, flashcards, or SOL reference."
)

if __name__ == "__main__":
    iface.launch()