"""
extractor.py — the pipeline.

  Stage 1: render_pdf_pages()      PDF -> high-DPI PNGs
  Stage 2: detect_regions()        page PNG -> regions.json
  Stage 3: group_into_articles()   regions.json -> articles.json
              Primary:  Claude API (spatial reasoning over regions)
              Fallback: column-aware + dateline-anchored Python rules
  Stage 4: crop_articles()         articles.json -> article PNGs
  Stage 5: ocr_articles()          article PNGs -> article TXTs
"""
import os
import re
import json
import base64
import shutil
import threading
from io import BytesIO
from pathlib import Path

try:
    from generate_pdf import generate_political_pdf
except ImportError:
    generate_political_pdf = None

# Load .env file if it exists, to populate environment variables like ANTHROPIC_API_KEY
if not os.environ.get("ANTHROPIC_API_KEY"):
    env_path = Path(__file__).resolve().parent / ".env"
    if env_path.exists():
        try:
            for line in env_path.read_text(encoding="utf-8").splitlines():
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    key, val = line.split("=", 1)
                    if key.strip() == "ANTHROPIC_API_KEY":
                        os.environ["ANTHROPIC_API_KEY"] = val.strip().strip('"').strip("'")
        except Exception as e:
            print(f"Error loading .env file: {e}")

import fitz  # PyMuPDF
from PIL import Image

JOBS = {}
JOB_LOCK = threading.Lock()

_LAYOUT = None
_LAYOUT_LOCK = threading.Lock()


def _set_status(job_id, msg):
    with JOB_LOCK:
        if job_id in JOBS:
            JOBS[job_id]["message"] = msg


# ---------------------------------------------------------------------------
# Stage 1 — render PDF pages
# ---------------------------------------------------------------------------
def render_pdf_pages(pdf_path, out_dir, target_w, target_h, max_pages=None, pages_to_process=None):
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    doc = fitz.open(pdf_path)
    pages = []
    
    # Determine page selection target set
    targets = set(pages_to_process) if pages_to_process is not None else None

    for i, page in enumerate(doc, start=1):
        if max_pages and i > max_pages:
            break
        if targets is not None and i not in targets:
            continue
        rect = page.rect
        scale = min(target_w / rect.width, target_h / rect.height)
        mat = fitz.Matrix(scale, scale)
        pix = page.get_pixmap(matrix=mat, alpha=False)
        out_path = out_dir / f"page_{i:03d}.png"
        pix.save(str(out_path))
        pages.append({"page": i, "path": str(out_path),
                      "width": pix.width, "height": pix.height})
    doc.close()
    return pages


# ---------------------------------------------------------------------------
# Stage 2 — detect regions with Surya Layout
# ---------------------------------------------------------------------------
_SURYA_LAYOUT_MODEL = None
_SURYA_LAYOUT_PROCESSOR = None

def _load_layout_model():
    global _SURYA_LAYOUT_MODEL, _SURYA_LAYOUT_PROCESSOR
    if _SURYA_LAYOUT_MODEL is not None:
        return _SURYA_LAYOUT_MODEL, _SURYA_LAYOUT_PROCESSOR
    with _LAYOUT_LOCK:
        if _SURYA_LAYOUT_MODEL is not None:
            return _SURYA_LAYOUT_MODEL, _SURYA_LAYOUT_PROCESSOR
        import surya.settings
        from surya.model.detection.model import load_model, load_processor
        
        checkpoint = surya.settings.settings.LAYOUT_MODEL_CHECKPOINT
        _SURYA_LAYOUT_MODEL = load_model(checkpoint=checkpoint)
        _SURYA_LAYOUT_PROCESSOR = load_processor(checkpoint=checkpoint)
        return _SURYA_LAYOUT_MODEL, _SURYA_LAYOUT_PROCESSOR


def detect_regions(page_png_path):
    model, processor = _load_layout_model()
    from surya.layout import batch_layout_detection
    img = Image.open(page_png_path).convert("RGB")

    orig_w, orig_h = img.size
    scale_factor = 1.0
    # Surya has a numpy int32 overflow bug on Windows for very large images.
    # Downscale for layout detection (Surya internally downscales anyway).
    if max(orig_w, orig_h) > 2048:
        scale_factor = 2048.0 / max(orig_w, orig_h)
        new_w = int(orig_w * scale_factor)
        new_h = int(orig_h * scale_factor)
        img_for_surya = img.resize((new_w, new_h), Image.LANCZOS)
    else:
        img_for_surya = img

    layout_predictions = batch_layout_detection([img_for_surya], model, processor)

    regions = []

    # Label mapping from Surya to our expected types
    label_map = {
        "Picture": "image",
        "Section-header": "paragraph_title",
        "Title": "doc_title",
        "Text": "text",
        "Caption": "text",
        "Table": "table",
        "Page-header": "header",
        "Page-footer": "footer",
        "List-item": "list"
    }

    if layout_predictions:
        layout = layout_predictions[0]
        if hasattr(layout, 'bboxes'):
            for rid, box in enumerate(layout.bboxes, start=1):
                # Scale coordinates back up to original image dimensions
                bbox = [int(v / scale_factor) for v in box.bbox]
                raw_label = box.label
                label = label_map.get(raw_label, "text").lower()
                # Surya doesn't expose confidence in this API level easily, default to 1.0
                regions.append({"id": rid, "type": label, "bbox": bbox, "confidence": 1.0})

    # Detailed logging
    print(f"\n{'='*70}")
    print(f"LAYOUT DETECTION RESULTS: {len(regions)} regions found")
    print(f"{'='*70}")

    # Count by type
    type_counts = {}
    for r in regions:
        t = r["type"]
        type_counts[t] = type_counts.get(t, 0) + 1
    print(f"Region types: {dict(sorted(type_counts.items(), key=lambda x: -x[1]))}")
    print()

    # Print each region
    for r in regions:
        x1, y1, x2, y2 = r["bbox"]
        w = x2 - x1
        h = y2 - y1
        print(f"  Region {r['id']:3d}: type={r['type']:15s} bbox=[{x1:5d},{y1:5d},{x2:5d},{y2:5d}] size={w:5d}x{h:5d}px conf={r['confidence']:.2f}")
    print(f"{'='*70}\n")

    # Draw bounding boxes on page image and save as debug image
    try:
        from PIL import ImageDraw, ImageFont
        debug_img = Image.open(page_png_path).copy()
        draw = ImageDraw.Draw(debug_img)

        # Color map for different region types
        colors = {
            "text": "blue", "image": "green", "title": "red", "doc_title": "red",
            "header": "gray", "footer": "gray", "paragraph_title": "orange",
            "table": "purple", "figure": "green", "list": "cyan",
        }

        for r in regions:
            x1, y1, x2, y2 = r["bbox"]
            color = colors.get(r["type"], "yellow")
            # Draw rectangle
            draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
            # Draw label
            label_text = f"{r['id']}:{r['type']}"
            draw.text((x1 + 5, y1 + 5), label_text, fill=color)

        # Save debug image next to the page image
        debug_path = page_png_path.replace(".png", "_regions_debug.png")
        debug_img.save(debug_path)
        print(f"DEBUG IMAGE saved: {debug_path}")
        print(f"  Open this image to see all {len(regions)} bounding boxes drawn on the page")
        print(f"  Colors: text=blue, title/doc_title=red, image/figure=green, header/footer=gray, paragraph_title=orange")
        print()
    except Exception as e:
        print(f"  Could not create debug image: {e}")

    return regions


# ---------------------------------------------------------------------------
# Stage 3 — group regions into articles
# ---------------------------------------------------------------------------
def group_into_articles(page_png_path, regions, page_width, page_img=None):
    """Geometric grouping. Returns (articles, enriched_regions)."""
    # Filter out page-level chrome to prevent unneccessary top/bottom content
    valid_regions = [r for r in regions if r["type"] not in ("header", "footer", "page_number")]
    
    print("[running geometric column grouping]")
    import geometric_grouper
    articles = geometric_grouper.group_surya_regions_geometrically(valid_regions)
    return articles, valid_regions


# ---------------------------------------------------------------------------
# Stage 4 — crop each article
# ---------------------------------------------------------------------------
def crop_headlines(page_png_path, articles, regions, out_dir, page_num):
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    img = Image.open(page_png_path)
    
    region_by_id = {r["id"]: r for r in regions}
    results = []
    
    for art in articles:
        # Find the headline box
        title_id = art.get("title_region")
        if title_id and title_id in region_by_id:
            l, t, r, b = region_by_id[title_id]["bbox"]
        else:
            # Fallback to the top 15% of the article if no title region
            l, t, r, b = art["bbox"]
            b = min(b, int(t + (b - t) * 0.15))
            
        # Pad slightly
        l = max(0, l - 10)
        t = max(0, t - 10)
        r = min(img.width, r + 10)
        b = min(img.height, b + 10)
        
        crop = img.crop((l, t, r, b))
        art_name = f"p{page_num:03d}_a{art['article_id']:03d}"
        art_dir = out_dir / art_name
        art_dir.mkdir(parents=True, exist_ok=True)
        img_path = art_dir / "headline.png"
        crop.save(img_path)
        
        results.append({
            "article_id": art["article_id"],
            "name": art_name,
            "dir": str(art_dir),
            "headline_img_path": str(img_path),
            "bbox": art["bbox"],
            "title_region": title_id,
            "member_regions": art.get("member_regions", [])
        })
    return results


def crop_full_articles(page_png_path, selected_ids, articles, out_dir, page_num):
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    img = Image.open(page_png_path)

    results = []
    for art in articles:
        art_name = f"p{page_num:03d}_a{art['article_id']:03d}"
        if art_name not in selected_ids:
            continue
            
        l, t, r, b = art["bbox"]
        
        # Dynamic Padding: expand by up to 25px, but don't intersect other articles
        pad_l = pad_t = pad_r = pad_b = 25
        
        for other in articles:
            if other["article_id"] == art["article_id"]:
                continue
            ol, ot, or_, ob = other["bbox"]
            
            # Check vertical overlap
            if max(t, ot) < min(b, ob):
                if ol >= r:  # Other is to the right
                    pad_r = min(pad_r, max(0, (ol - r) // 2))
                elif or_ <= l: # Other is to the left
                    pad_l = min(pad_l, max(0, (l - or_) // 2))
                    
            # Check horizontal overlap
            if max(l, ol) < min(r, or_):
                if ot >= b:  # Other is below
                    pad_b = min(pad_b, max(0, (ot - b) // 2))
                elif ob <= t:  # Other is above
                    pad_t = min(pad_t, max(0, (t - ob) // 2))

        l = max(0, int(l - pad_l))
        t = max(0, int(t - pad_t))
        r = min(img.width, int(r + pad_r))
        b = min(img.height, int(b + pad_b))

        crop = img.crop((l, t, r, b))
        name = f"p{page_num:03d}_a{art['article_id']:03d}"
        art_dir = out_dir / name
        art_dir.mkdir(parents=True, exist_ok=True)
        crop.save(art_dir / "article.png")

        info = {
            "page": page_num,
            "article_id": art["article_id"],
            "bbox": [l, t, r, b],
            "member_regions": art["member_regions"],
            "title_region": art.get("title_region"),
            "dateline": art.get("dateline"),
            "headline_hint": art.get("headline_hint"),
            "grouped_by": art.get("grouped_by", "rules"),
        }
        (art_dir / "info.json").write_text(json.dumps(info, indent=2, ensure_ascii=False), encoding="utf-8")
        results.append({"name": name, "dir": str(art_dir), **info})
    return results


# ---------------------------------------------------------------------------
# Stage 5 — OCR each cropped article (Telugu)
# ---------------------------------------------------------------------------
_paddle_ocr_instance = None

def _ocr_with_paddleocr(img_path):
    """Use PaddleOCR for better reading-order-aware Telugu OCR."""
    global _paddle_ocr_instance
    temp_path = None
    try:
        from PIL import Image
        with Image.open(img_path) as im:
            px = im.width * im.height
            if px > 4000000:
                # Instead of skipping, resize the article image to prevent PaddleOCR deadlock
                ratio = (3800000 / px) ** 0.5
                new_w = int(im.width * ratio)
                new_h = int(im.height * ratio)
                print(f"Article image {img_path.name} is too large ({im.width}x{im.height} = {px}px). Resizing to {new_w}x{new_h} for safe PaddleOCR.")
                resized_im = im.resize((new_w, new_h), Image.LANCZOS)
                temp_path = img_path.parent / f"temp_resized_{img_path.name}"
                resized_im.save(temp_path)
        
        import paddleocr
        if _paddle_ocr_instance is None:
            _paddle_ocr_instance = paddleocr.PaddleOCR(lang='te')
        
        path_to_ocr = str(temp_path) if temp_path else str(img_path)
        result = _paddle_ocr_instance.predict(path_to_ocr)

        lines = []
        if isinstance(result, list) and len(result) > 0:
            det_result = result[0]
            # PaddleOCR v3.5 returns DetResult with 'rec_texts' or nested structure
            rec_texts = None
            if hasattr(det_result, 'get'):
                rec_texts = det_result.get('rec_texts', None)
            if rec_texts:
                lines = list(rec_texts)
            else:
                # Try to extract from boxes/text pairs
                boxes = det_result.get('dt_polys', []) if hasattr(det_result, 'get') else []
                texts = det_result.get('rec_texts', []) if hasattr(det_result, 'get') else []
                if texts:
                    # Sort by y-coordinate (top to bottom), then x (left to right)
                    # to get proper reading order
                    pairs = []
                    for i, txt in enumerate(texts):
                        if i < len(boxes):
                            box = boxes[i]
                            if hasattr(box, 'tolist'):
                                box = box.tolist()
                            # Get top-left y coordinate for sorting
                            if isinstance(box, (list, tuple)) and len(box) >= 4:
                                y = min(box[j] for j in range(1, len(box), 2)) if len(box) >= 8 else box[1]
                                x = min(box[j] for j in range(0, len(box), 2)) if len(box) >= 8 else box[0]
                            else:
                                y, x = 0, 0
                            pairs.append((y, x, txt))
                        else:
                            pairs.append((0, 0, txt))

                    # Sort: primarily by y (row), then x (column within row)
                    # Group into rows based on y proximity
                    if pairs:
                        pairs.sort(key=lambda p: (p[0], p[1]))
                        row_threshold = 20  # pixels
                        rows = []
                        current_row = [pairs[0]]
                        for p in pairs[1:]:
                            if abs(p[0] - current_row[0][0]) < row_threshold:
                                current_row.append(p)
                            else:
                                current_row.sort(key=lambda p: p[1])  # sort by x within row
                                rows.append(current_row)
                                current_row = [p]
                        current_row.sort(key=lambda p: p[1])
                        rows.append(current_row)

                        for row in rows:
                            lines.append(" ".join(p[2] for p in row))
                else:
                    # Last resort: try str representation
                    try:
                        s = str(det_result)
                        if len(s) > 10:
                            return s
                    except:
                        pass

        return "\n".join(lines) if lines else None
    except Exception as e:
        print(f"PaddleOCR text extraction failed: {e}")
        return None
    finally:
        if temp_path and temp_path.exists():
            try:
                temp_path.unlink()
            except:
                pass


def ocr_headlines(headline_records):
    for rec in headline_records:
        img_path = Path(rec["headline_img_path"])
        if not img_path.exists():
            rec["headline_text"] = ""
            continue
        text = _ocr_with_paddleocr(img_path)
        rec["headline_text"] = (text or "").strip()
        (img_path.parent / "headline.txt").write_text(rec["headline_text"], encoding="utf-8")

def ocr_articles(article_records):
    for rec in article_records:
        art_dir = Path(rec["dir"])
        img_path = art_dir / "article.png"
        if not img_path.exists():
            continue

        text = _ocr_with_paddleocr(img_path)
        (art_dir / "article.txt").write_text(text or "", encoding="utf-8")

        info_path = art_dir / "info.json"
        if info_path.exists():
            info = json.loads(info_path.read_text(encoding="utf-8"))
            info_path.write_text(json.dumps(info, indent=2, ensure_ascii=False), encoding="utf-8")


# ---------------------------------------------------------------------------
# Political Filter — crop first, then filter with Claude
# ---------------------------------------------------------------------------
def _triage_headlines_with_claude(headline_records, page_num, job_dir=None):
    """Send HEADLINES ONLY to Claude to select political articles, saving 90% compute."""
    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        # If no key, select all to avoid losing data
        return [r["name"] for r in headline_records]

    # Build the list of text articles
    article_list_data = []
    for rec in headline_records:
        txt = rec.get("headline_text", "").strip()
        if not txt:
            txt = "[No text extracted by OCR]"
        
        article_list_data.append({
            "id": rec["name"],
            "headline_text": txt
        })

    if not article_list_data:
        return []

    try:
        import anthropic
        client = anthropic.Anthropic(api_key=api_key)

        articles_json_str = json.dumps(article_list_data, indent=2, ensure_ascii=False)

        prompt = f"""You are a senior opposition party strategist in Telangana state, India.
I have extracted the HEADLINES from page {page_num} of a Telugu newspaper. Here is the list:

{articles_json_str}

Return the IDs of articles that MIGHT describe DAMAGE, LOSS, CORRUPTION, or SUFFERING happening to people in Telangana, where the government could be blamed.
Also return the IDs of any articles about farmer protests, NEET paper leaks, or protests against the ruling party.
If you are unsure, INCLUDE it. We just want to filter out obvious ads, sports, and unrelated filler.

Return ONLY a valid JSON list of IDs. Example: ["p001_a002", "p001_a005"]. No prose."""

        response = client.messages.create(
            model="claude-sonnet-4-6",
            max_tokens=1000,
            messages=[{"role": "user", "content": prompt}],
        )

        text = response.content[0].text.strip()
        if text.startswith("```"):
            text = text.split("\n", 1)[1]
            if text.endswith("```"):
                text = text[:-3]

        selected_ids = json.loads(text.strip())
        print(f"  Page {page_num}: Triage kept {len(selected_ids)}/{len(headline_records)} articles based on headlines.")
        return selected_ids

    except Exception as e:
        print(f"Headline triage failed: {e}")
        return [r["name"] for r in headline_records]


def _analyze_articles_with_claude(records, page_num, job_dir=None):
    """Deep analysis on the full OCR text of the selected articles."""
    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        return None

    article_list_data = []
    for rec in records:
        art_dir = Path(rec["dir"])
        txt_path = art_dir / "article.txt"
        txt_content = ""
        if txt_path.exists():
            txt_content = txt_path.read_text(encoding="utf-8").strip()
        if not txt_content:
            txt_content = "[No text extracted by OCR]"
            
        article_list_data.append({
            "id": rec["name"],
            "ocr_text": txt_content[:3000]
        })

    if not article_list_data:
        return []

    try:
        import anthropic
        client = anthropic.Anthropic(api_key=api_key)

        # Convert list to JSON string for the prompt
        articles_json_str = json.dumps(article_list_data, indent=2, ensure_ascii=False)

        prompt = f"""You are a senior opposition party strategist in Telangana state, India.

I have run OCR on the cropped articles of page {page_num} of a Telugu newspaper. Here is the list of article IDs and their Telugu OCR text contents:

{articles_json_str}

For EACH article, ask yourself these TWO questions:

QUESTION 1: "Does this article describe actual DAMAGE, LOSS, DEATH, CORRUPTION, or SUFFERING happening to people in Telangana? And can the government (state OR central) be BLAMED for it — either for causing it, failing to prevent it, or failing to respond to it?"
- The article must describe a NEGATIVE EVENT that has already occurred — not a future plan, not an announcement, not a statistic.
- There must be a VICTIM — someone who died, lost money, was cheated, is suffering, or was harmed.
- The government (Telangana state OR central/BJP) can be blamed — for causing it, failing to prevent it, or failing to respond to it.
- Central government failures also qualify IF the article describes actual harm to Telangana citizens (e.g., NEET paper leak harming Telangana students, central policy causing job losses in Telangana).
- If the article describes something POSITIVE (new scheme, development, achievement) -> answer is NO.
- If the article describes a PLAN, PETITION, or ANNOUNCEMENT (e.g. MLA/civilians submitting a representation requesting funds, work commencement, inspections, or temple matters) -> answer is NO. It must describe an active issue/damage event, not requests for future work.
- If there is NO victim and NO damage described -> answer is NO.
- RESPONSIBILITY CHECK: Ask "WHO caused this damage?" If the damage was caused by nature (fire, flood, lightning), individual negligence (reckless driving, personal accident), criminals (theft, murder), or private parties — and the government had NO role in causing or preventing it — answer is NO. Only answer YES if the government is CLEARLY responsible through its policy, negligence, corruption, or failure to act.
- If YES -> proceed to Question 2
- If NO -> REJECT this article immediately

QUESTION 2: "Can the opposition party DIRECTLY USE this article to attack the Telangana government in a press conference, assembly session, or public rally?"
- If YES -> SELECT this article
- If NO -> REJECT this article

BOTH answers must be YES to select. If either answer is NO, reject.

DEDUPLICATION: If two or more articles cover the SAME issue or event (e.g., same farmer protest, same accident, same scam), only SELECT the ONE with the most complete coverage or the larger article. REJECT the duplicates. Do not keep two articles about the same news story.

NOW go through EACH article ID I gave you. Ask Q1 and Q2 for each. Only select if BOTH are YES.

IMPORTANT: This application is for TELANGANA STATE politics only. Farmer issues are the MOST IMPORTANT — always select if farmers are suffering.

DO NOT SELECT — automatically REJECT:
- International news (wars, foreign affairs, immigration)
- Government achievements, new schemes, welfare announcements, development projects
- Newspaper masthead, headers, advertisements, classifieds
- Sports, entertainment, cinema, horoscopes
- General statistics without actual damage (birth rate, population, GDP)
- News about other states
- Resignations — ALWAYS REJECT unless the article explicitly contains words like corruption, scam, fired, forced out, or political pressure. A person resigning for personal, health, family, or workload reasons is NOT a government failure.
- PM Modi speeches or party meetings (unless specific policy failure hurting Telangana people)
- Individual accidents (road accidents, electric shocks, drownings, fire accidents) — these are NOT government failures, they can happen due to personal negligence or bad luck. Only select if there is a MASS disaster (10+ deaths) or CLEAR systemic government negligence (e.g., bridge collapse, building collapse, factory fire due to no safety inspections).

Return a JSON object with TWO lists:

1. "reasoning": for EVERY article (selected or not), explain your thinking (KEEP ALL REASONS EXTREMELY BRIEF, MAX 10 WORDS EACH TO PREVENT TOKEN CUTOFF):
  - "id": the article ID
  - "headline_english": brief English headline (max 8 words)
  - "q1": "YES" or "NO"
  - "q1_reason": short reason (max 10 words)
  - "q2": "YES" or "NO" (skip if q1 is NO)
  - "q2_reason": short reason (max 10 words, skip if q1 is NO)
  - "decision": "SELECT" or "REJECT"

2. "selected": only the articles where both Q1 and Q2 are YES:
  - "id": the article ID
  - "headline_english": English translation of the headline
  - "q1_answer": answer to Question 1 (1 short sentence)
  - "q2_answer": answer to Question 2 (1 short sentence)
  - "attack_angle": what opposition should say (1 sentence)
  - "priority": "high" or "medium"
  - "category": one of corruption, governance_failure, public_grievance, law_order, health, education, infrastructure, ruling_party_crisis, farmer_issues

Return ONLY valid JSON. No markdown. If no articles pass both questions, return {{"reasoning": [...], "selected": []}}"""

        response = client.messages.create(
            model="claude-sonnet-4-6",
            max_tokens=8192,
            messages=[{"role": "user", "content": prompt}],
        )

        text = response.content[0].text.strip()
        if text.startswith("```"):
            text = text.split("\n", 1)[1]
            if text.endswith("```"):
                text = text[:-3]

        result = json.loads(text.strip())

        # Log full reasoning to file and terminal
        reasoning = result.get("reasoning", [])
        if reasoning:
            print(f"\n=== POLITICAL FILTER REASONING (Page {page_num}) ===")
            for r in reasoning:
                status = "✅ SELECT" if r.get("decision") == "SELECT" else "❌ REJECT"
                print(f"  {r.get('id','?')}: {r.get('headline_english','')[:50]}")
                print(f"    Q1: {r.get('q1','?')} — {r.get('q1_reason','')}")
                if r.get('q2'):
                    print(f"    Q2: {r.get('q2','?')} — {r.get('q2_reason','')}")
                print(f"    -> {status}")
            print(f"=== END REASONING ===\n")

        # Save reasoning log to job directory
        log_path = Path(job_dir) / f"filter_reasoning_page{page_num}.json" if job_dir else None
        if log_path:
            log_path.write_text(json.dumps(result, indent=2, ensure_ascii=False), encoding="utf-8")

        return result.get("selected", [])

    except Exception as e:
        print(f"Political filter failed: {e}")
        import traceback; traceback.print_exc()
        return None


def _generate_opposition_brief(selected_articles, job_dir):
    """Generate a readable opposition daily brief text file."""
    high = [a for a in selected_articles if a.get("priority") == "high"]
    medium = [a for a in selected_articles if a.get("priority") == "medium"]

    brief_path = Path(job_dir) / "opposition_daily_brief.txt"
    with open(brief_path, "w", encoding="utf-8") as f:
        f.write("=" * 70 + "\n")
        f.write("  OPPOSITION DAILY BRIEF\n")
        f.write("=" * 70 + "\n\n")
        f.write(f"Total significant articles: {len(selected_articles)}\n")
        f.write(f"HIGH: {len(high)} | MEDIUM: {len(medium)}\n")
        f.write("-" * 70 + "\n\n")

        for label, group in [("HIGH PRIORITY — IMMEDIATE ACTION", high),
                              ("MEDIUM PRIORITY", medium)]:
            if group:
                f.write(f">>> {label} <<<\n\n")
                for i, a in enumerate(group, 1):
                    f.write(f"{i}. [{a.get('id', '?')}] {a.get('headline_english', 'N/A')}\n")
                    f.write(f"   Category: {a.get('category', '?')}\n")
                    f.write(f"   WHY: {a.get('political_significance', 'N/A')}\n")
                    f.write(f"   ATTACK: {a.get('attack_angle', 'N/A')}\n\n")
                f.write("-" * 70 + "\n")

        f.write("END OF BRIEF\n")


# ---------------------------------------------------------------------------
# Orchestrator
# ---------------------------------------------------------------------------
def process_pdf(pdf_path, job_dir, target_w, target_h, job_id, max_pages=None, pages_to_process=None):
    job_dir = Path(job_dir)
    pages_dir = job_dir / "pages"
    articles_root = job_dir / "articles"
    pages_dir.mkdir(parents=True, exist_ok=True)
    articles_root.mkdir(parents=True, exist_ok=True)

    # Step 1: Render PDF pages
    _set_status(job_id, "Rendering PDF pages...")
    pages = render_pdf_pages(pdf_path, pages_dir, target_w, target_h, max_pages=max_pages, pages_to_process=pages_to_process)

    all_records = []
    all_selected = []

    for p in pages:
        n = p["page"]
        page_img_path = p["path"]

        # Step 2: PaddleOCR layout detection — precise bounding boxes
        _set_status(job_id, f"Page {n}: detecting layout regions...")
        regions = detect_regions(page_img_path)
        (pages_dir / f"page_{n:03d}.regions.json").write_text(
            json.dumps({"page": n, "regions": regions}, indent=2, ensure_ascii=False), encoding="utf-8"
        )
        print(f"  Page {n}: {len(regions)} layout regions detected")

        # Step 3: Group regions into articles
        _set_status(job_id, f"Page {n}: grouping articles...")
        page_img = Image.open(page_img_path)
        articles, enriched = group_into_articles(page_img_path, regions, p["width"], page_img=page_img)

        (pages_dir / f"page_{n:03d}.articles.json").write_text(
            json.dumps({"page": n, "articles": articles, "enriched_regions": enriched},
                       indent=2, ensure_ascii=False), encoding="utf-8"
        )

        # --- DETAILED GROUPING LOG ---
        print(f"\n{'='*70}")
        print(f"GROUPING RESULTS — PAGE {n}")
        print(f"{'='*70}")
        print(f"  Input: {len(regions)} regions -> Output: {len(articles)} articles")
        print()

        # Build region lookup
        region_map = {r["id"]: r for r in regions}

        for art in articles:
            aid = art.get("article_id", "?")
            member_ids = art.get("member_region_ids", art.get("member_regions", []))
            title_id = art.get("title_region_id", art.get("title_region"))
            hint = (art.get("headline_hint") or "")[:60]
            dateline = art.get("dateline", "")
            grouped_by = art.get("grouped_by", "")
            bbox = art.get("bbox", [])

            print(f"  ARTICLE {aid}: \"{hint}\"")
            if dateline:
                print(f"    Dateline: {dateline}")
            if title_id:
                title_r = region_map.get(title_id, {})
                print(f"    Title region: ID {title_id} (type={title_r.get('type','?')}, bbox={title_r.get('bbox',[])})")
            print(f"    Member regions ({len(member_ids)}): {member_ids}")

            # Show each member region's details
            if member_ids:
                min_x, min_y, max_x, max_y = 99999, 99999, 0, 0
                for rid in member_ids:
                    r = region_map.get(rid, {})
                    rb = r.get("bbox", [0,0,0,0])
                    print(f"      Region {rid:3d}: type={r.get('type','?'):15s} bbox=[{rb[0]:5d},{rb[1]:5d},{rb[2]:5d},{rb[3]:5d}] size={rb[2]-rb[0]:5d}x{rb[3]-rb[1]:5d}px")
                    min_x = min(min_x, rb[0])
                    min_y = min(min_y, rb[1])
                    max_x = max(max_x, rb[2])
                    max_y = max(max_y, rb[3])
                print(f"    Merged bbox: [{min_x},{min_y},{max_x},{max_y}] size={max_x-min_x}x{max_y-min_y}px")
            if bbox:
                print(f"    Article bbox: {bbox}")
            print()

        # Summary stats
        region_counts = [len(art.get("member_region_ids", art.get("member_regions", []))) for art in articles]
        if region_counts:
            print(f"  STATS:")
            print(f"    Articles: {len(articles)}")
            print(f"    Regions per article: min={min(region_counts)}, max={max(region_counts)}, avg={sum(region_counts)/len(region_counts):.1f}")
            large = [i+1 for i,c in enumerate(region_counts) if c > 6]
            if large:
                print(f"    ⚠️ Large articles (>6 regions): {large} — may have merged multiple stories")
        print(f"{'='*70}\n")

        # Draw grouped articles on debug image with different colors per article
        try:
            from PIL import ImageDraw
            debug_img = Image.open(page_img_path).copy()
            draw = ImageDraw.Draw(debug_img)

            # Dark distinct colors for each article
            article_colors = [
                (255, 0, 0),       # red
                (0, 0, 255),       # blue
                (0, 150, 0),       # dark green
                (180, 0, 180),     # purple
                (200, 100, 0),     # dark orange
                (0, 150, 150),     # teal
                (150, 0, 0),       # dark red
                (0, 0, 150),       # dark blue
                (100, 100, 0),     # olive
                (200, 0, 100),     # magenta
                (0, 100, 200),     # steel blue
                (150, 75, 0),      # brown
                (100, 0, 150),     # dark purple
                (0, 130, 80),      # dark teal
                (180, 50, 50),     # brick red
                (50, 50, 180),     # indigo
                (0, 180, 0),       # green
                (180, 0, 0),       # crimson
                (0, 0, 180),       # navy
                (150, 150, 0),     # dark yellow
                (200, 50, 150),    # pink
                (50, 150, 50),     # forest green
                (100, 50, 200),    # violet
                (200, 150, 50),    # gold
                (50, 100, 150),    # slate
            ]

            for idx, art in enumerate(articles):
                color = article_colors[idx % len(article_colors)]
                aid = art.get("article_id", idx+1)
                member_ids = art.get("member_region_ids", art.get("member_regions", []))

                for rid in member_ids:
                    r = region_map.get(rid, {})
                    rb = r.get("bbox", [0,0,0,0])
                    # Draw thick border with article color
                    draw.rectangle([rb[0], rb[1], rb[2], rb[3]], outline=color, width=5)
                    # Label with article ID
                    draw.text((rb[0]+8, rb[1]+8), f"A{aid}", fill=color)

                # Draw merged article bbox as thick dark outer border
                if member_ids:
                    all_bboxes = [region_map.get(rid, {}).get("bbox", [0,0,0,0]) for rid in member_ids]
                    mx1 = min(b[0] for b in all_bboxes)
                    my1 = min(b[1] for b in all_bboxes)
                    mx2 = max(b[2] for b in all_bboxes)
                    my2 = max(b[3] for b in all_bboxes)
                    # Thick dark outer border around the whole article
                    draw.rectangle([mx1-15, my1-15, mx2+15, my2+15], outline=color, width=15)
                    # Dark background label
                    label = f"Article {aid}: {(art.get('headline_hint') or '')[:30]}"
                    draw.rectangle([mx1-10, my1-40, mx1 + len(label)*12, my1-10], fill=color)
                    draw.text((mx1-5, my1-38), label, fill="white")

            debug_path = page_img_path.replace(".png", "_articles_debug.png")
            debug_img.save(debug_path)
            print(f"ARTICLES DEBUG IMAGE saved: {debug_path}")
            print(f"  Each article has a unique color. Regions with same color = same article.")
            print()
        except Exception as e:
            print(f"  Could not create articles debug image: {e}")
        # Step 4: Crop and OCR HEADLINES only
        _set_status(job_id, f"Page {n}: cropping {len(articles)} headlines...")
        headline_records = crop_headlines(page_img_path, articles, regions, articles_root, n)
        
        _set_status(job_id, f"Page {n}: running fast OCR on {len(headline_records)} headlines...")
        ocr_headlines(headline_records)

        # Step 5: Triage with Claude Haiku based on headlines
        _set_status(job_id, f"Page {n}: triaging headlines with Claude...")
        selected_ids = _triage_headlines_with_claude(headline_records, n, job_dir=str(job_dir))

        # Cleanup non-selected directories
        for rec in headline_records:
            if rec["name"] not in selected_ids:
                shutil.rmtree(rec["dir"], ignore_errors=True)

        if not selected_ids:
            print(f"  Page {n}: 0 political articles found during triage.")
            continue

        # Step 6: Crop FULL articles for the selected ones
        _set_status(job_id, f"Page {n}: cropping {len(selected_ids)} FULL articles...")
        records = crop_full_articles(page_img_path, selected_ids, articles, articles_root, n)
        print(f"  Page {n}: cropped {len(records)} full articles")

        # Step 7: OCR full articles
        _set_status(job_id, f"Page {n}: running heavy OCR on {len(records)} full articles...")
        ocr_articles(records)

        # Step 8: Deep Analysis with Claude Opus
        _set_status(job_id, f"Page {n}: deep political analysis with Claude...")
        selected = _analyze_articles_with_claude(records, n, job_dir=str(job_dir))

        if selected:
            final_ids = {s["id"] for s in selected}
            print(f"  Page {n}: {len(selected)} politically significant articles: {sorted(final_ids)}")

            # Keep only selected articles, remove the rest
            for rec in records:
                art_dir = Path(rec["dir"])
                art_name = art_dir.name
                if art_name in final_ids:
                    # Update info.json with political analysis
                    sel_info = next((s for s in selected if s["id"] == art_name), {})
                    info_path = art_dir / "info.json"
                    if info_path.exists():
                        info = json.loads(info_path.read_text(encoding="utf-8"))
                        info["headline_english"] = sel_info.get("headline_english", "")
                        info["political_significance"] = sel_info.get("political_significance", "")
                        info["attack_angle"] = sel_info.get("attack_angle", "")
                        info["priority"] = sel_info.get("priority", "")
                        info["category"] = sel_info.get("category", "")
                        info["grouped_by"] = "paddleocr + political_filter"
                        info_path.write_text(json.dumps(info, indent=2, ensure_ascii=False), encoding="utf-8")
                    all_records.append(rec)
                    sel_info["image_file"] = f"{art_name}.png"
                    all_selected.append(sel_info)
                else:
                    shutil.rmtree(str(art_dir), ignore_errors=True)
        else:
            # No API key or no political articles — keep all cropped articles
            print(f"  Page {n}: no final political filter applied, keeping all {len(records)} articles")
            all_records.extend(records)

    # --- SUMMARY ACROSS ALL PAGES ---
    print(f"\n{'='*70}")
    print(f"LAYOUT DETECTION SUMMARY — ALL PAGES")
    print(f"{'='*70}")

    # Reload all regions from saved JSON
    all_page_summaries = []
    grand_total = 0
    grand_type_counts = {}

    for p in pages:
        n = p["page"]
        regions_file = pages_dir / f"page_{n:03d}.regions.json"
        if regions_file.exists():
            data = json.loads(regions_file.read_text(encoding="utf-8"))
            page_regions = data.get("regions", [])
            grand_total += len(page_regions)

            # Count by type
            type_counts = {}
            for r in page_regions:
                t = r["type"]
                type_counts[t] = type_counts.get(t, 0) + 1
                grand_type_counts[t] = grand_type_counts.get(t, 0) + 1

            # Estimate articles: doc_title + paragraph_title = potential article starts
            doc_titles = type_counts.get("doc_title", 0)
            para_titles = type_counts.get("paragraph_title", 0)

            print(f"\n  PAGE {n}:")
            print(f"    Total regions: {len(page_regions)}")
            print(f"    Types: {dict(sorted(type_counts.items(), key=lambda x: -x[1]))}")
            print(f"    Potential articles (doc_title): {doc_titles}")
            print(f"    Potential sub-articles (paragraph_title): {para_titles}")
            print(f"    Estimated article range: {doc_titles} (min) to {doc_titles + para_titles} (max)")

            all_page_summaries.append({
                "page": n,
                "total_regions": len(page_regions),
                "type_counts": type_counts,
                "doc_titles": doc_titles,
                "paragraph_titles": para_titles,
            })

    print(f"\n  GRAND TOTAL:")
    print(f"    Pages: {len(pages)}")
    print(f"    Total regions: {grand_total}")
    print(f"    All types: {dict(sorted(grand_type_counts.items(), key=lambda x: -x[1]))}")
    total_doc = grand_type_counts.get("doc_title", 0)
    total_para = grand_type_counts.get("paragraph_title", 0)
    print(f"    Estimated articles across all pages: {total_doc} (min) to {total_doc + total_para} (max)")
    print(f"{'='*70}\n")

    # Save summary JSON
    summary_path = job_dir / "layout_summary.json"
    summary_path.write_text(json.dumps({
        "total_pages": len(pages),
        "total_regions": grand_total,
        "grand_type_counts": grand_type_counts,
        "estimated_articles_min": total_doc,
        "estimated_articles_max": total_doc + total_para,
        "pages": all_page_summaries,
    }, indent=2, ensure_ascii=False), encoding="utf-8")
    print(f"Summary saved: {summary_path}")

    # Save political filter results
    if all_selected:
        (job_dir / "political_articles.json").write_text(
            json.dumps(all_selected, indent=2, ensure_ascii=False), encoding="utf-8"
        )
        _generate_opposition_brief(all_selected, job_dir)

    # Copy all kept article images into one folder
    if all_records:
        all_images_dir = job_dir / "all_political_images"
        all_images_dir.mkdir(parents=True, exist_ok=True)
        for rec in all_records:
            art_dir = Path(rec["dir"])
            img_path = art_dir / "article.png"
            if img_path.exists():
                art_name = art_dir.name
                shutil.copy2(str(img_path), str(all_images_dir / f"{art_name}.png"))
        print(f"  Copied {len(all_records)} images to {all_images_dir}")

    # Generate the combined PDF with AI insights and images
    if all_selected and generate_political_pdf:
        generate_political_pdf(all_selected, job_dir)

    _set_status(job_id, "Done!")
    return {"pages": len(pages), "articles": len(all_records)}