import os
import gradio as gr
from pypdf import PdfReader

from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import PyPDFLoader
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_core.documents import Document

# -------------------------
# Config
# -------------------------
APP_TITLE = "Election Memory Engine"
ARCHIVE_DIR = "archive"  # put your PDFs inside this folder in the Space repo
MODEL_NAME = "gpt-4o-mini"

# -------------------------
# Document loading
# -------------------------
def list_pdfs(root_dir: str):
    pdfs = []
    for root, _, files in os.walk(root_dir):
        for f in files:
            if f.lower().endswith(".pdf"):
                pdfs.append(os.path.join(root, f))
    return sorted(pdfs)

def load_archive_docs():
    pdf_paths = list_pdfs(ARCHIVE_DIR)
    docs = []
    for path in pdf_paths:
        try:
            loader = PyPDFLoader(path)
            loaded = loader.load()
            # Ensure each page knows which file it came from
            for d in loaded:
                d.metadata = d.metadata or {}
                d.metadata["source"] = os.path.basename(path)
            docs.extend(loaded)
        except Exception:
            # Fallback: try simple reader if loader fails
            try:
                reader = PdfReader(path)
                text = "\n".join([(p.extract_text() or "") for p in reader.pages])
                docs.append(Document(page_content=text, metadata={"source": os.path.basename(path)}))
            except Exception:
                continue
    return docs

# -------------------------
# Index build (startup)
# -------------------------
def build_vectorstore():
    docs = load_archive_docs()
    if not docs:
        return None, "No PDFs found. Upload PDFs into the `archive/` folder in Files."

    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
    chunks = splitter.split_documents(docs)

    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
    vs = FAISS.from_documents(chunks, embeddings)

    return vs, f"✅ Indexed {len(docs)} pages/sections from {len(set([d.metadata.get('source','') for d in docs]))} file(s)."

VECTORSTORE, INDEX_STATUS = build_vectorstore()

# -------------------------
# Answering with sources
# -------------------------
def answer_question(question: str):
    if VECTORSTORE is None:
        return INDEX_STATUS

    if not question or not question.strip():
        return "Type a question first."

    retriever = VECTORSTORE.as_retriever(search_kwargs={"k": 5})
    hits = retriever.get_relevant_documents(question)

    # Build a compact context for the model + collect citations
    context_parts = []
    cite_lines = []
    for i, d in enumerate(hits, start=1):
        src = d.metadata.get("source", "unknown")
        excerpt = " ".join(d.page_content.split())[:260]
        cite_lines.append(f"{i}. **{src}** — {excerpt}…")
        context_parts.append(f"[Source {i}: {src}]\n{d.page_content}")

    llm = ChatOpenAI(model=MODEL_NAME, temperature=0.2)

    prompt = f"""
You are an archive assistant for an election newsroom.
Answer using ONLY the context below. If the context is insufficient, say what is missing and suggest what document would help.
Write in a newsroom-ready format:
- 5–8 bullet background points
- If dates exist, include a short timeline
- End with "Sources used: 1,2,3..."

Question: {question}

Context:
{chr(10).join(context_parts)}
""".strip()

    resp = llm.invoke(prompt).content
    citations = "\n".join(cite_lines) if cite_lines else "_No sources returned._"

    return f"{resp}\n\n---\n### Sources (snippets)\n{citations}"

# -------------------------
# UI
# -------------------------
CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:wght@400;600;700&display=swap');

.gradio-container{
  background: white !important;
  font-family: 'IBM Plex Mono', monospace !important;
}

h1{
  font-size:52px !important;
  font-weight:800 !important;
  line-height:1.05 !important;
  letter-spacing:-1px;
}

h3{
  color:#6b7280 !important;
  font-weight:400 !important;
}

button{
  border-radius: 12px !important;
  padding: 12px 24px !important;
  font-weight: 600 !important;
}

.gradio-container > div:first-child {
  border-bottom: 1px solid #e5e7eb;
  padding-bottom: 20px;
  margin-bottom: 26px;
}
"""

LOGO_PATH = "assets/archive logo.png"
HERO_PATH = "assets/file cabinet.png"

def load_pil(path: str):
    try:
        from PIL import Image
        if os.path.exists(path):
            return Image.open(path)
    except Exception:
        return None
    return None

with gr.Blocks(css=CUSTOM_CSS, title=APP_TITLE) as demo:

    # Navbar
    with gr.Row():
        with gr.Column(scale=1, min_width=70):
            logo = load_pil(LOGO_PATH)
            if logo is not None:
                gr.Image(value=logo, show_label=False, container=False, height=44)
            else:
                gr.Markdown("<span style='color:#6b7280;font-size:12px;'>Logo missing</span>")

        with gr.Column(scale=5):
            gr.Markdown("<div style='font-weight:700; font-size:18px; padding-top:6px;'>Election Archive</div>")

        with gr.Column(scale=6):
            gr.HTML(
                """
                <div style="display:flex; justify-content:flex-end; gap:28px; font-weight:500; padding-top:10px;">
                  <span>Home</span>
                  <span>Archive</span>
                  <span>Upload</span>
                </div>
                """
            )

    # Hero
    with gr.Row():
        with gr.Column(scale=6):
            gr.Markdown("""
# What Happened  
# Last  
# Election?

### Turn your newsroom archive into a searchable memory engine.
### Upload documents. Ask questions. Get answers with sources.
""")
            gr.Markdown(f"<div style='color:#6b7280; font-size:12px;'>Index status: {INDEX_STATUS}</div>")

        with gr.Column(scale=6):
            hero = load_pil(HERO_PATH)
            if hero is not None:
                gr.Image(value=hero, show_label=False, container=False)
            else:
                gr.Markdown("<span style='color:#6b7280;font-size:12px;'>Hero image missing</span>")

    # Search / Q&A
    with gr.Row():
        with gr.Column():
            q = gr.Textbox(
                label="Ask the archive",
                placeholder="e.g. What were the major issues in 2022, and what patterns repeat from 2017?",
            )
            out = gr.Markdown()
            btn = gr.Button("Search", variant="primary")

    btn.click(answer_question, q, out)
    q.submit(answer_question, q, out)

demo.launch()