import os
import faiss
import gradio as gr
import numpy as np
import requests

from typing import List
from PyPDF2 import PdfReader
from docx import Document
from sentence_transformers import SentenceTransformer

# Set Groq API key (set this as a secret in Hugging Face Spaces)
GROQ_API_KEY = os.getenv("GROQ_API_KEY")

# Load SentenceTransformer model
embed_model = SentenceTransformer("all-MiniLM-L6-v2")

# Global storage
index = None
text_chunks = []

# System prompt for Groq LLaMA3
SYSTEM_PROMPT = """
You are an AI study supervisor. Answer the student's question using only the content from the uploaded document.
Be concise, supportive, and accurate. If the answer is not in the document, say:
"I could not find this information in your document."
"""

def extract_text_from_pdf(file_path):
    reader = PdfReader(file_path)
    return " ".join([page.extract_text() or "" for page in reader.pages])

def extract_text_from_docx(file_path):
    doc = Document(file_path)
    return "\n".join([para.text for para in doc.paragraphs])

def file_to_text_chunks(file) -> List[str]:
    """Reads PDF or DOCX file and splits text into chunks"""
    ext = os.path.splitext(file.name)[1].lower()
    if ext == ".pdf":
        raw_text = extract_text_from_pdf(file)
    elif ext == ".docx":
        raw_text = extract_text_from_docx(file)
    else:
        raise ValueError("Unsupported file type. Please upload a PDF or DOCX.")
    
    # Split into chunks
    chunk_size = 500
    return [raw_text[i:i + chunk_size] for i in range(0, len(raw_text), chunk_size)]

def embed_chunks(chunks: List[str]) -> faiss.IndexFlatL2:
    """Embed and store chunks in FAISS"""
    embeddings = embed_model.encode(chunks)
    index = faiss.IndexFlatL2(embeddings.shape[1])
    index.add(np.array(embeddings))
    return index

def retrieve_context(question: str, k: int = 3) -> str:
    """Find relevant chunks for a question"""
    question_vec = embed_model.encode([question])
    distances, indices = index.search(np.array(question_vec), k)
    return "\n".join([text_chunks[i] for i in indices[0]])

def ask_llama3(question: str, context: str) -> str:
    """Query Groq's LLaMA3 API"""
    url = "https://api.groq.com/openai/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {GROQ_API_KEY}",
        "Content-Type": "application/json"
    }

    payload = {
        "model": "llama-3.3-70b-versatile",
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT.strip()},
            {"role": "user", "content": f"Answer the following question using the provided document content:\n{context}\n\nQuestion: {question}"}
        ],
        "temperature": 0.3
    }

    response = requests.post(url, headers=headers, json=payload)
    if response.status_code == 200:
        return response.json()["choices"][0]["message"]["content"]
    else:
        return f"⚠️ Error contacting Groq API: {response.status_code} - {response.text}"

def process_file(file):
    global index, text_chunks
    try:
        text_chunks = file_to_text_chunks(file)
        index = embed_chunks(text_chunks)
        return "✅ Document uploaded and processed. You can now ask questions."
    except Exception as e:
        return f"❌ Error: {str(e)}"

def answer_question(user_question):
    if not index:
        return "❌ Please upload a document first."
    context = retrieve_context(user_question)
    return ask_llama3(user_question, context)

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("🧠 **AI Study Supervisor**\nUpload a study document (PDF or DOCX) and ask questions about it.")
    
    with gr.Row():
        file_input = gr.File(label="📄 Upload PDF or DOCX", file_types=[".pdf", ".docx"])
        upload_status = gr.Textbox(label="Status", interactive=False)

    file_input.change(fn=process_file, inputs=file_input, outputs=upload_status)
    
    with gr.Row():
        user_input = gr.Textbox(label="💬 Ask a Question")
        submit_btn = gr.Button("Submit")
    
    answer_output = gr.Textbox(label="🧠 Answer")
    submit_btn.click(fn=answer_question, inputs=user_input, outputs=answer_output)

demo.launch()