import os
import re
import gradio as gr
from langchain.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from langchain.vectorstores import Chroma
from langchain.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=embedding_model,
    persist_directory="geometry_db",  # relative folder inside your Hugging Face Space
    collection_name="geometry_sol"
)

# Load OpenAI key (you must add 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)

# Prompt templates
templates = {
    "general": PromptTemplate(
        input_variables=["context", "query"],
        template="""
You are a strict assistant for the Virginia Geometry SOL.

Only use exact phrases from the following SOL text:
{context}

Answer the question: "{query}"

If the answer is in the SOL text, quote it exactly. Do not rephrase or summarize. Do not add your own explanation.

If the answer is not in the context, reply: "The answer is not found in the provided SOL text."
"""
    ),
    "lesson plan": PromptTemplate(
        input_variables=["context", "query"],
        template="""
Given the following retrieved SOL text:
{context}

Generate a Geometry lesson plan based on: "{query}"
Include:
1. Simple explanation of the concept.
2. Real-world example.
3. Engaging class activity.
Be concise and curriculum-aligned for high school.
"""
    ),
    "worksheet": PromptTemplate(
        input_variables=["context", "query"],
        template="""
{context}

Create a student worksheet for: "{query}"
Include: concept summary, a worked example, and 3 practice problems.
"""
    ),
    "proofs": PromptTemplate(
        input_variables=["context", "query"],
        template="""
{context}

Generate a proof-focused geometry lesson plan for: "{query}"
Include: student-friendly explanation, real-world link, and activity.
"""
    )
}

# 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
def rag_query(query, mode="general"):
    docs = vector_store.similarity_search(query, k=2)
    context = "\n\n".join([doc.page_content for doc in docs])
    prompt = templates[mode].format_prompt(context=context, query=query).to_string()
    return llm.invoke(prompt).content

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

# Gradio UI
iface = gr.Interface(
    fn=ask_geometry_sol,
    inputs=[
        gr.Textbox(label="Enter your Geometry SOL question or topic"),
        gr.Radio(["general", "lesson plan", "worksheet", "proofs"], value="general", label="Response type")
    ],
    outputs="text",
    title="📘 Virginia Geometry SOL Assistant",
    description="Ask about any 2023 Geometry SOL (Standards of Learning). Get exact quotes, lesson plans, worksheets, or proof-based lessons."
)

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