import os
import zipfile
import gradio as gr
from langchain_openai import ChatOpenAI
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain

# Unzip vector DB if not already extracted
if not os.path.exists("geometry_chroma"):
    with zipfile.ZipFile("geometry_chroma.zip", 'r') as zip_ref:
        zip_ref.extractall(".")

# Load vector DB
embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vectordb = Chroma(persist_directory="geometry_chroma", embedding_function=embedding_model)
retriever = vectordb.as_retriever()

# Set OpenAI key (use Secrets or .env later)
os.environ["OPENAI_API_KEY"] = os.getenv("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")

llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.2)

# Prompt templates
templates = {
    "flashcard": PromptTemplate(
        input_variables=["context", "query"],
        template="""
{context}

Create 5 flashcards based on: "{query}"
Each card should have a question and short answer.
"""
    )
}

def generate_output(prompt_type, query):
    docs = retriever.get_relevant_documents(query)
    context = "\n\n".join([doc.page_content for doc in docs])
    chain = LLMChain(llm=llm, prompt=templates[prompt_type])
    return chain.run({"context": context, "query": query})

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 📐 Geometry Assistant")

    query = gr.Textbox(label="Enter a geometry topic")
    prompt_type = gr.Dropdown(["flashcard"], value="flashcard", label="Prompt Type")
    output = gr.Textbox(label="Generated Output", lines=12)

    btn = gr.Button("Generate")
    btn.click(fn=generate_output, inputs=[prompt_type, query], outputs=output)

demo.launch()