import os
import tempfile
import fitz  # PyMuPDF
import streamlit as st
import torch
from transformers import AutoConfig, AutoTokenizer, T5ForConditionalGeneration

# Page Configuration
st.set_page_config(
    page_title="RTI & Civic PDF Simplifier",
    page_icon="📄",
    layout="centered",
)

st.title("📄 Civic & RTI PDF Simplifier")
st.markdown(
    "Upload bureaucratic, legal, or RTI PDF documents to generate concise,"
    " plain-language bullet points."
)

MODEL_REPO = "CaseLoop/pdf"
FALLBACK_BASE = "google/flan-t5-base"


@st.cache_resource(show_spinner="Loading model weights...")
def load_model_and_tokenizer():
  # Tokenizer fallback
  try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, legacy=False)
  except Exception:
    tokenizer = AutoTokenizer.from_pretrained(FALLBACK_BASE, legacy=False)

  # Model loader with architecture fallback
  try:
    model = T5ForConditionalGeneration.from_pretrained(MODEL_REPO)
  except Exception:
    try:
      config = AutoConfig.from_pretrained(FALLBACK_BASE)
      model = T5ForConditionalGeneration.from_pretrained(
          MODEL_REPO, config=config
      )
    except Exception:
      model = T5ForConditionalGeneration.from_pretrained(FALLBACK_BASE)

  model.eval()
  return tokenizer, model


tokenizer, model = load_model_and_tokenizer()


def extract_chunks(
    pdf_path: str, chunk_size: int = 180, max_chunks: int = 6
) -> list[str]:
  """Extracts lines from a PDF and chunks words for FLAN-T5 input limits."""
  doc = fitz.open(pdf_path)
  lines = []
  for page in doc:
    text = page.get_text("text")
    lines.extend([
        line.strip()
        for line in text.split("\n")
        if len(line.strip().split()) >= 5
    ])

  words = " ".join(lines).split()
  if not words:
    return []

  chunks = [
      " ".join(words[i : i + chunk_size])
      for i in range(0, len(words), chunk_size)
  ]
  return chunks[:max_chunks]


# UI File Uploader
uploaded_file = st.file_uploader(
    "Choose a PDF file", type=["pdf"], help="Upload an RTI reply or civic notice"
)

if uploaded_file is not None:
  if st.button("Simplify Document", type="primary"):
    with st.spinner("Parsing document and running simplification..."):
      with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
        tmp.write(uploaded_file.read())
        tmp_path = tmp.name

      chunks = extract_chunks(tmp_path)
      if os.path.exists(tmp_path):
        os.remove(tmp_path)

      if not chunks:
        st.warning(
            "No extractable text found in this PDF. It might be a scanned image."
        )
      else:
        st.subheader("Key Takeaways in Plain Language:")
        for idx, chunk in enumerate(chunks, 1):
          prompt = f"simplify bureaucratic text into simple language: {chunk}"
          inputs = tokenizer(
              prompt, return_tensors="pt", max_length=512, truncation=True
          )

          with torch.no_grad():
            output_tokens = model.generate(
                **inputs,
                max_length=120,
                min_length=20,
                num_beams=4,
                no_repeat_ngram_size=3,
                repetition_penalty=2.5,
                early_stopping=True,
            )

          summary = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
          st.markdown(f"**{idx}.** {summary}")