import os
import gradio as gr
import numpy as np
import faiss
from pypdf import PdfReader
from sentence_transformers import SentenceTransformer
from groq import Groq

# ---------------------------
# CONFIG
# ---------------------------
GROQ_API_KEY = os.getenv("GROQ_API_KEY")

if not GROQ_API_KEY:
    raise ValueError("Missing GROQ_API_KEY in Hugging Face Secrets")

client = Groq(api_key=GROQ_API_KEY)

# ---------------------------
# EMBEDDING MODEL
# ---------------------------
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")

dimension = 384
index = faiss.IndexFlatL2(dimension)
stored_chunks = []

# ---------------------------
# PDF LOADER
# ---------------------------
def load_pdf(file):
    reader = PdfReader(file)
    text = ""

    for page in reader.pages:
        page_text = page.extract_text()
        if page_text:
            text += page_text + "\n"

    return text

# ---------------------------
# CHUNKING
# ---------------------------
def chunk_text(text, chunk_size=800, overlap=150):
    chunks = []
    start = 0

    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        start = end - overlap

    return chunks

# ---------------------------
# BUILD VECTOR DB
# ---------------------------
def build_index(text):
    global stored_chunks, index

    stored_chunks = chunk_text(text)

    embeddings = embedding_model.encode(stored_chunks)
    embeddings = np.array(embeddings).astype("float32")

    index = faiss.IndexFlatL2(dimension)
    index.add(embeddings)

# ---------------------------
# RETRIEVAL
# ---------------------------
def retrieve(query, k=4):
    query_vec = embedding_model.encode([query]).astype("float32")

    distances, indices = index.search(query_vec, k)

    results = []
    for i in indices[0]:
        if i < len(stored_chunks):
            results.append(stored_chunks[i])

    return results

# ---------------------------
# LLM (GROQ)
# ---------------------------
def ask_llm(context, question):

    prompt = f"""
You are InsightPilot AI — a senior AI research analyst at a top consulting firm.

Your job is to analyze documents and produce decision-grade insights.

Context:
{context}

Question:
{question}

Return structured output:

1. Executive Summary
2. Key Insights
3. Risks & Limitations
4. Opportunities / Actions
5. Final Recommendation (score 0-10 with reasoning)
6. Confidence Level (high/medium/low with explanation)
"""

    response = client.chat.completions.create(
        model="llama3-70b-8192",
        messages=[{"role": "user", "content": prompt}]
    )

    return response.choices[0].message.content

# ---------------------------
# PIPELINE: CHAT
# ---------------------------
def chat_with_doc(file, question):
    if file is None:
        return "⚠️ Please upload a PDF first."

    text = load_pdf(file)

    if not text.strip():
        return "⚠️ Could not extract text from PDF."

    build_index(text)

    docs = retrieve(question)
    context = "\n".join(docs)

    return ask_llm(context, question)

# ---------------------------
# PIPELINE: REPORT
# ---------------------------
def generate_report(file):
    if file is None:
        return "⚠️ Please upload a PDF first."

    text = load_pdf(file)
    build_index(text)

    docs = retrieve("summary risks opportunities insights")
    context = "\n".join(docs)

    prompt = f"""
You are InsightPilot AI — a senior consulting analyst.

Create a structured strategic report:

1. Executive Summary
2. Key Insights
3. Risks
4. Opportunities
5. Strategic Recommendation (0-10 score)
6. Confidence Level

Context:
{context}
"""

    response = client.chat.completions.create(
        model="llama3-70b-8192",
        messages=[{"role": "user", "content": prompt}]
    )

    return response.choices[0].message.content

# ---------------------------
# UI (SAAS-GRADE)
# ---------------------------
with gr.Blocks(theme=gr.themes.Soft()) as app:

    gr.Markdown("""
# 🧠 InsightPilot AI
### AI-Powered Document Intelligence & Decision Engine

Turn any PDF into structured, decision-ready insights like a top consultant.

---

### 🚀 What it does:
- Extracts intelligence from documents
- Identifies risks & opportunities
- Generates strategic recommendations
""")

    with gr.Tab("📄 Chat with Document"):
        file_input = gr.File(label="Upload PDF")
        question = gr.Textbox(
            label="Ask a question",
            placeholder="Example: What are the main risks in this document?"
        )
        output = gr.Textbox(label="AI Analysis", lines=15)

        btn = gr.Button("Generate Insight")

        btn.click(
            fn=chat_with_doc,
            inputs=[file_input, question],
            outputs=output
        )

    with gr.Tab("📊 Generate Strategic Report"):
        file_input2 = gr.File(label="Upload PDF")
        output2 = gr.Textbox(label="Strategic Report", lines=20)

        btn2 = gr.Button("Generate Report")

        btn2.click(
            fn=generate_report,
            inputs=file_input2,
            outputs=output2
        )

    gr.Markdown("""
---
### ⚡ InsightPilot AI
Built with RAG + Groq + FAISS for decision intelligence systems
""")

app.launch()