# app.py

import streamlit as st
import fitz  # PyMuPDF
import datetime
import json
from sentence_transformers import SentenceTransformer, util

# === Load embedding model ===
@st.cache_resource
def load_model():
    return SentenceTransformer('all-MiniLM-L6-v2')

model = load_model()

# === Extract blocks from PDF ===
def extract_blocks(pdf_path):
    doc = fitz.open(pdf_path)
    blocks = []
    for page_num, page in enumerate(doc, 1):
        for block in page.get_text("dict")["blocks"]:
            if "lines" not in block:
                continue
            text = " ".join(span["text"] for line in block["lines"] for span in line["spans"]).strip()
            if text:
                blocks.append({
                    "text": text,
                    "page": page_num
                })
    return blocks

# === Rank blocks using similarity to task ===
def rank_blocks(blocks, persona, job, source):
    task_prompt = f"{persona} - {job}"
    task_embed = model.encode(task_prompt, convert_to_tensor=True)

    results = []
    for block in blocks:
        block_embed = model.encode(block["text"], convert_to_tensor=True)
        sim = util.cos_sim(task_embed, block_embed).item()
        results.append({
            "text": block["text"],
            "page": block["page"],
            "score": sim,
            "source": source
        })

    return sorted(results, key=lambda x: x["score"], reverse=True)

# === Build final output JSON ===
def build_output(ranked_blocks, input_files, persona, job):
    top_blocks = ranked_blocks[:5]
    output = {
        "metadata": {
            "input_documents": input_files,
            "persona": persona,
            "job_to_be_done": job,
            "processing_timestamp": str(datetime.datetime.now())
        },
        "extracted_sections": [],
        "subsection_analysis": []
    }

    for i, block in enumerate(top_blocks):
        output["extracted_sections"].append({
            "document": block["source"],
            "section_title": block["text"][:50],
            "importance_rank": i + 1,
            "page_number": block["page"]
        })
        output["subsection_analysis"].append({
            "document": block["source"],
            "refined_text": block["text"],
            "page_number": block["page"]
        })

    return output

# === Streamlit Interface ===
st.title("📘 Adobe GenSolve 1B: Semantic Section Extractor")

persona = st.text_input("🧑 Persona", value="Travel Planner")
job = st.text_input("🎯 Job to be done", value="Plan a trip of 4 days for a group of 10 college friends.")

uploaded_files = st.file_uploader("📄 Upload PDF files", type="pdf", accept_multiple_files=True)

if st.button("🚀 Extract Sections") and uploaded_files:
    all_ranked = []
    filenames = []

    for f in uploaded_files:
        path = f.name
        with open(path, "wb") as out_file:
            out_file.write(f.read())
        filenames.append(path)

        blocks = extract_blocks(path)
        ranked = rank_blocks(blocks, persona, job, path)
        all_ranked.extend(ranked)

    output = build_output(all_ranked, filenames, persona, job)

    # Show and download
    st.success("✅ Extraction Complete")
    st.json(output)

    with open("output.json", "w") as f:
        json.dump(output, f, indent=2)

    st.download_button("📥 Download output.json", data=json.dumps(output, indent=2), file_name="output.json", mime="application/json")