import os
import subprocess
import gradio as gr
from langchain_community.vectorstores import Weaviate
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from sentence_transformers import SentenceTransformer
import weaviate

# Groq import
from langchain_groq import ChatGroq 

# === SECRETS ===
WEAVIATE_URL = os.environ["WEAVIATE_URL"]
WEAVIATE_KEY = os.environ["WEAVIATE_KEY"]
os.environ["GROQ_API_KEY"] = os.environ["GROQ_API_KEY"]

# AUTO-INGEST ON START
def run_ingestion():
    print("Running ingestion...")
    result = subprocess.run(["python", "ingest.py"], capture_output=True, text=True)
    print(result.stdout)
    if result.returncode != 0:
        print("Ingestion failed:", result.stderr)
    else:
        print("Ingestion complete!")

# Run once at startup
run_ingestion()

# RAG CHAIN
@gr.cache
def get_rag_chain():
    # v3 Weaviate Client
    client = weaviate.Client(
        url=WEAVIATE_URL,
        auth_client_secret=weaviate.AuthApiKey(WEAVIATE_KEY)
    )
    embedder = SentenceTransformer("all-MiniLM-L6-v2")
    vectorstore = Weaviate(client, "Paper", "text", embedding=embedder)
    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

    # Use ChatGroq
    llm = ChatGroq(model="llama-3.1-70b-instruct", temperature=0)

    prompt = PromptTemplate.from_template(
        "Answer using only this context:\n{context}\n\nQuestion: {question}\nAnswer:"
    )
    return RetrievalQA.from_chain_type(
        llm=llm,
        retriever=retriever,
        chain_type_kwargs={"prompt": prompt}
    )

def search(query):
    try:
        qa = get_rag_chain()
        result = qa.invoke({"query": query})
        answer = result["result"]

        sources = "\n\n".join([
            f"**{doc.metadata.get('title', 'No Title')}**\n{doc.page_content[:300]}..."
            for doc in result["source_documents"]
        ])
        return answer, sources if sources.strip() else "No sources retrieved."
    except Exception as e:
        return f"Error: {str(e)}", "Check logs."

# GRADIO UI
with gr.Blocks(title="ArXiv RAG Search") as demo:
    gr.Markdown("# ArXiv RAG Search Engine")
    gr.Markdown("10K+ research papers • Llama-3.1 • Weaviate • Auto-indexed")

    with gr.Row():
        txt = gr.Textbox(
            label="Ask about AI, physics, math...",
            placeholder="What is the attention mechanism?",
            lines=2
        )
        btn = gr.Button("Search", variant="primary")

    with gr.Row():
        out_answer = gr.Markdown(label="Answer")
        out_sources = gr.Markdown(label="Sources")

    btn.click(search, inputs=txt, outputs=[out_answer, out_sources])
    txt.submit(lambda x: (x, gr.update()), txt, [txt, btn]).then(search, txt, [out_answer, out_sources])

demo.launch()