import spaces  # MUST be the absolute first import to prevent ZeroGPU CUDA initialization crash
import fitz  # PyMuPDF
import gradio as gr
import torch
from transformers import AutoConfig, AutoTokenizer, T5ForConditionalGeneration
import re
import json

# ============================================================
# MODEL 1: PDF SIMPLIFIER
# ============================================================
SIMPLIFIER_REPO = "Ronak0/last-caseloop-pdf-simplifier"
FALLBACK_BASE = "google/flan-t5-base"

try:
  simplifier_tokenizer = AutoTokenizer.from_pretrained(SIMPLIFIER_REPO, legacy=False)
except Exception:
  simplifier_tokenizer = AutoTokenizer.from_pretrained(FALLBACK_BASE, legacy=False)

try:
  simplifier_model = T5ForConditionalGeneration.from_pretrained(SIMPLIFIER_REPO)
except Exception:
  try:
    config = AutoConfig.from_pretrained(FALLBACK_BASE)
    simplifier_model = T5ForConditionalGeneration.from_pretrained(SIMPLIFIER_REPO, config=config)
  except Exception:
    simplifier_model = T5ForConditionalGeneration.from_pretrained(FALLBACK_BASE)


def extract_chunks(
    pdf_path: str, chunk_size: int = 180, max_chunks: int = 6
) -> list[str]:
  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()) >= 4
    ])

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


def clean_summary_output(text: str) -> str:
  """Splits clumped bullet points onto separate lines and cleans formatting."""
  normalized = re.sub(r'\s*•\s*', '\n• ', text).strip()
  if not normalized.startswith('•') and not re.match(r'^\d+\.', normalized):
    normalized = '• ' + normalized

  lines = [line.strip() for line in normalized.split('\n') if line.strip()]
  cleaned_lines = []
  for line in lines:
    line = re.sub(r'\b(\w+)\s+(.*?)\b\1\b', r'\1 \2', line, flags=re.IGNORECASE)
    cleaned_lines.append(line)

  return '\n'.join(cleaned_lines)


@spaces.GPU
def simplify_document(pdf_file):
  if pdf_file is None:
    return "Please upload a valid PDF file."

  chunks = extract_chunks(pdf_file.name)
  if not chunks:
    return (
        "No readable text found in document. It may contain scanned image pages."
    )

  device = "cuda" if torch.cuda.is_available() else "cpu"
  simplifier_model.to(device)

  takeaways = []
  for idx, chunk in enumerate(chunks, 1):
    prompt = f"Simplify civic/RTI document:\n{chunk}"

    inputs = simplifier_tokenizer(
        prompt, return_tensors="pt", max_length=512, truncation=True
    ).to(device)

    with torch.no_grad():
      output_tokens = simplifier_model.generate(
          **inputs,
          max_new_tokens=256,
          min_length=20,
          num_beams=4,
          no_repeat_ngram_size=3,
          repetition_penalty=1.15,
          length_penalty=0.9,
          early_stopping=True,
      )

    summary = simplifier_tokenizer.decode(output_tokens[0], skip_special_tokens=True)
    summary = clean_summary_output(summary)

    takeaways.append(f"Section {idx}:\n{summary}")

  return "\n\n".join(takeaways)


# ============================================================
# MODEL 2: RIGHTS NAVIGATOR
# ============================================================
NAVIGATOR_REPO = "CaseLoop/rights-navigator"

try:
  navigator_tokenizer = AutoTokenizer.from_pretrained(NAVIGATOR_REPO, legacy=False)
except Exception:
  navigator_tokenizer = simplifier_tokenizer

try:
  navigator_model = T5ForConditionalGeneration.from_pretrained(NAVIGATOR_REPO)
except Exception:
  navigator_model = simplifier_model


@spaces.GPU
def navigate_rights(grievance_text, jurisdiction):
  if not grievance_text or not grievance_text.strip():
    return "Please describe your grievance or dispute."

  device = "cuda" if torch.cuda.is_available() else "cpu"
  navigator_model.to(device)

  prompt = f'[INST] Analyze citizen grievance: "{grievance_text}" in jurisdiction "{jurisdiction}". Return strict JSON with keys: domain, statutory_hooks, legal_rights, immediate_actions, escalation_pathway. [/INST]'

  inputs = navigator_tokenizer(
      prompt, return_tensors="pt", max_length=512, truncation=True
  ).to(device)

  with torch.no_grad():
    output_tokens = navigator_model.generate(
        **inputs,
        max_new_tokens=300,
        temperature=0.2,
        num_beams=2,
    )

  gen_text = navigator_tokenizer.decode(output_tokens[0], skip_special_tokens=True)

  try:
    match = re.search(r'\{.*\}', gen_text, re.DOTALL)
    if match:
      data = json.loads(match.group(0))
      return f"""### Domain: {data.get('domain', 'Civic Dispute')}
**Statutory Hooks:** {', '.join(data.get('statutory_hooks', []))}

**1. Your Legal Rights:**
{chr(10).join([f'- {r}' for r in data.get('legal_rights', [])])}

**2. Immediate Actionable Steps:**
{chr(10).join([f'{i+1}. {a}' for i, a in enumerate(data.get('immediate_actions', []))])}

**3. Where to Escalate:**
{chr(10).join([f'- {e}' for e in data.get('escalation_pathway', [])])}"""
  except Exception:
    pass

  return f"""### Domain: Consumer / Civil Grievance
**Statutory Hooks:** Consumer Protection Framework & Public Grievance Norms

**1. Your Legal Rights:**
- Right to formal written justification for delayed or withheld statutory fulfillment.
- Right to seek financial and operational remedy under applicable acts.

**2. Immediate Actionable Steps:**
1. Compile all communication logs, notices, agreements, and payment transaction receipts.
2. Issue a formal pre-litigation notice granting standard statutory resolution timeline.

**3. Where to Escalate:**
- Lodge formal complaint via national e-portal or regulatory commission in {jurisdiction}.
- File petition before the corresponding District Redressal Forum."""


# ============================================================
# INTERFACES
# ============================================================
pdf_tab = gr.Interface(
    fn=simplify_document,
    inputs=gr.File(label="Upload Civic / RTI PDF", file_types=[".pdf"]),
    outputs=gr.Textbox(
        label="Plain-Language Takeaways", lines=12, interactive=False
    ),
    title="📄 Civic & RTI PDF Simplifier",
    description=(
        "Upload complex bureaucratic documents to extract simplified summaries"
        " using fine-tuned FLAN-T5."
    ),
    flagging_mode="never",
)

rights_tab = gr.Interface(
    fn=navigate_rights,
    inputs=[
        gr.Textbox(label="Describe Your Grievance / Dispute", lines=4, placeholder="e.g., Landlord refused to return security deposit after 30 days..."),
        gr.Textbox(label="Jurisdiction / City", value="Delhi (NCR)")
    ],
    outputs=gr.Markdown(label="Actionable Legal Guidance"),
    title="⚖️ Civic Rights Navigator",
    description="Analyze dispute facts to extract legal rights, immediate steps, and escalation pathways.",
    flagging_mode="never",
)

demo = gr.TabbedInterface(
    [pdf_tab, rights_tab],
    ["📄 PDF Simplifier", "⚖️ Rights Navigator"]
)

if __name__ == "__main__":
  demo.launch()