"""
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 threading
from io import BytesIO
from pathlib import Path

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):
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    doc = fitz.open(pdf_path)
    pages = []
    for i, page in enumerate(doc, start=1):
        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 PaddleOCR PP-Structure
# ---------------------------------------------------------------------------
def _load_layout_model():
    global _LAYOUT
    if _LAYOUT is not None:
        return _LAYOUT
    with _LAYOUT_LOCK:
        if _LAYOUT is not None:
            return _LAYOUT
        import paddleocr
        # PaddleOCR v3.5.0 — use LayoutDetection for layout analysis.
        # Stock ~0.5 confidence cutoff (no threshold override) — this is what
        # Claude grouping is tuned to see.
        try:
            _LAYOUT = paddleocr.LayoutDetection()
        except Exception:
            try:
                _LAYOUT = paddleocr.PPStructureV3()
            except Exception:
                _LAYOUT = paddleocr.PaddleOCR(lang='te')
        return _LAYOUT


def detect_regions(page_png_path):
    layout = _load_layout_model()
    import numpy as np
    img = np.array(Image.open(page_png_path).convert("RGB"))

    # PaddleOCR v3.5.0 uses .predict() instead of direct call
    if hasattr(layout, 'predict'):
        raw = layout.predict(page_png_path)
    elif callable(layout):
        raw = layout(img)
    else:
        raw = layout.predict(img)

    regions = []

    # PaddleOCR v3.5.0: raw is a list of DetResult objects (dict-like).
    # Each DetResult has 'boxes' key containing list of detected regions.
    # Each box has: 'cls_id', 'label', 'score', 'coordinate' (4 floats: x1,y1,x2,y2)
    if isinstance(raw, list) and len(raw) > 0:
        for det_result in raw:
            # Get boxes from the DetResult
            boxes = None
            if hasattr(det_result, 'get'):
                boxes = det_result.get('boxes', None)
            elif hasattr(det_result, 'boxes'):
                boxes = det_result.boxes
            elif isinstance(det_result, dict) and 'boxes' in det_result:
                boxes = det_result['boxes']

            if boxes is not None:
                for box in boxes:
                    rid = len(regions) + 1
                    # Extract coordinate
                    coord = None
                    if isinstance(box, dict):
                        coord = box.get('coordinate', box.get('bbox', None))
                        label = box.get('label', box.get('type', 'unknown')).lower()
                        score = float(box.get('score', box.get('confidence', 0)))
                    elif hasattr(box, 'coordinate'):
                        coord = box.coordinate
                        label = getattr(box, 'label', 'unknown').lower()
                        score = float(getattr(box, 'score', 0))
                    else:
                        label = 'unknown'
                        score = 0.0

                    bbox = [0, 0, 0, 0]
                    if coord is not None:
                        if hasattr(coord, 'tolist'):
                            coord = coord.tolist()
                        coord = [float(v) for v in coord]
                        if len(coord) == 4:
                            bbox = [int(v) for v in coord]
                        elif len(coord) >= 8:
                            xs = [coord[j] for j in range(0, len(coord), 2)]
                            ys = [coord[j] for j in range(1, len(coord), 2)]
                            bbox = [int(min(xs)), int(min(ys)), int(max(xs)), int(max(ys))]

                    regions.append({"id": rid, "type": label, "bbox": bbox, "confidence": score})

    # 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
#
# Primary: Claude API (set ANTHROPIC_API_KEY env var).
# Fallback: column-aware + dateline Python rules.
# ---------------------------------------------------------------------------

DATELINE_RE = re.compile(
    r"^[\u0C00-\u0C7F\u0964\u0965A-Za-z0-9\s.\-]{1,40}?,"
    r"\s*[\u0C00-\u0C7F]{1,8}\s*\d{1,2}"
    r"[^:]{0,30}\(ఆంధ్రజ్యోతి\)[\s\u0C00-\u0C7F]{0,4}:",
    re.UNICODE,
)


def looks_like_dateline(text):
    if not text:
        return None
    snippet = text.strip()[:120]
    m = DATELINE_RE.match(snippet)
    if m:
        return snippet[:m.end()]
    if "ఆంధ్రజ్యోతి" in snippet[:100] and ":" in snippet[:120]:
        return snippet[:snippet.find(":") + 1]
    return None


def _bbox_union(a, b):
    return [min(a[0], b[0]), min(a[1], b[1]), max(a[2], b[2]), max(a[3], b[3])]


def _make_topband_ocr(page_img):
    try:
        import pytesseract
    except ImportError:
        return None

    def topband(reg):
        l, t, r, b = reg["bbox"]
        strip_h = min(140, b - t)
        if strip_h <= 0:
            return ""
        strip = page_img.crop((l, t, r, t + strip_h))
        try:
            return pytesseract.image_to_string(strip, lang="tel")
        except Exception:
            return ""
    return topband


def _enrich_regions_with_text(regions, topband_ocr):
    enriched = []
    for r in regions:
        rec = dict(r)
        if topband_ocr and r["type"] in ("text", "paragraph", "list"):
            try:
                txt = topband_ocr(r) or ""
            except Exception:
                txt = ""
            rec["top_text"] = txt.strip()[:160]
            d = looks_like_dateline(rec["top_text"])
            rec["dateline"] = d if d else None
        else:
            rec["top_text"] = ""
            rec["dateline"] = None
        enriched.append(rec)
    return enriched


# ---------- Claude-based grouping ----------
def _make_page_thumbnail(page_png_path, max_dim=900):
    img = Image.open(page_png_path)
    img.thumbnail((max_dim, max_dim), Image.LANCZOS)
    buf = BytesIO()
    img.save(buf, format="JPEG", quality=80)
    return base64.standard_b64encode(buf.getvalue()).decode("ascii")


CLAUDE_GROUPING_PROMPT = """You group newspaper page regions into INDIVIDUAL articles. Each news story is a SEPARATE article.

You receive:
1. A thumbnail of the full page (for spatial context).
2. A JSON list of regions on the page. Each has:
   - id (int)
   - type ("title", "text", "figure", "table", etc.)
   - bbox [left, top, right, bottom] in pixels (top-left origin)
   - top_text: a Tesseract OCR snippet from the top of text regions
   - dateline: if non-null, the region starts an article body (strong signal)

CRITICAL RULES — BE CONSERVATIVE, SPLIT MORE:
1. A `title` or `doc_title` region ALWAYS starts a NEW article. Every title = separate article.
2. A region with a non-null `dateline` ALWAYS starts a new article body.
3. Text regions attach to the NEAREST title above them IN THE SAME COLUMN ONLY.
4. Newspapers have COLUMNS. A region in column 1 NEVER belongs to an article in column 2, even if vertically close.
5. When in doubt, SPLIT into separate articles rather than merge. Too many small articles is better than one giant article containing multiple stories.
6. An article should typically have 1-5 regions. If an article has more than 6 regions, you are probably merging multiple stories — split them.
7. A banner title spanning multiple columns owns ONLY the text directly below it, not the entire page.
8. Types "header", "footer", "page_number" — make these their own separate article.
9. Figures/images attach to the nearest title above them in THEIR column only.
10. NEVER merge two articles that have different headlines or topics into one.

Return ONLY valid JSON, no prose, no markdown fences:
{
  "articles": [
    {
      "article_id": 1,
      "title_region_id": <int or null>,
      "member_region_ids": [<int>, ...],
      "headline_hint": "<short string or null>",
      "dateline": "<string or null>"
    }
  ]
}

Order articles top-to-bottom, then left-to-right by topmost member region.
Every region (except headers/footers) appears in exactly one article.
AIM FOR 15-25 articles per page for a typical Telugu newspaper front page."""


def _group_with_claude(page_png_path, regions, api_key, model="claude-opus-4-8"):
    try:
        import anthropic
    except ImportError:
        raise RuntimeError("anthropic package not installed — pip install anthropic")

    client = anthropic.Anthropic(api_key=api_key)
    thumbnail_b64 = _make_page_thumbnail(page_png_path)

    compact = [
        {
            "id": r["id"],
            "type": r["type"],
            "bbox": r["bbox"],
            "top_text": r.get("top_text", "")[:160],
            "dateline": r.get("dateline"),
        }
        for r in regions
    ]

    user_content = [
        {"type": "image",
         "source": {"type": "base64", "media_type": "image/jpeg", "data": thumbnail_b64}},
        {"type": "text",
         "text": "Regions:\n" + json.dumps(compact, ensure_ascii=False, indent=2)},
    ]

    msg = client.messages.create(
        model=model,
        max_tokens=4000,
        system=CLAUDE_GROUPING_PROMPT,
        messages=[{"role": "user", "content": user_content}],
    )

    raw = "".join(b.text for b in msg.content if hasattr(b, "text")).strip()
    if raw.startswith("```"):
        raw = raw.split("\n", 1)[1] if "\n" in raw else raw
        if raw.endswith("```"):
            raw = raw[:raw.rfind("```")]
    data = json.loads(raw)

    region_by_id = {r["id"]: r for r in regions}
    articles = []
    for art in data.get("articles", []):
        ids = art.get("member_region_ids", [])
        if not ids:
            continue
        bbox = None
        for rid in ids:
            r = region_by_id.get(rid)
            if not r:
                continue
            bbox = list(r["bbox"]) if bbox is None else _bbox_union(bbox, r["bbox"])
        if bbox is None:
            continue
        articles.append({
            "article_id": art.get("article_id", len(articles) + 1),
            "title_region": art.get("title_region_id"),
            "member_regions": ids,
            "bbox": bbox,
            "dateline": art.get("dateline"),
            "headline_hint": art.get("headline_hint"),
            "grouped_by": "claude",
        })
    return articles


# ---------- Fallback Python heuristic ----------
def _assign_columns(regions, page_width):
    if not regions:
        return {0: regions}
    threshold = page_width * 0.04
    lefts = sorted({r["bbox"][0] for r in regions})
    column_starts = [lefts[0]]
    for prev, curr in zip(lefts, lefts[1:]):
        if curr - prev > threshold:
            column_starts.append(curr)
    columns = {i: [] for i in range(len(column_starts))}
    for r in regions:
        left = r["bbox"][0]
        col_idx = 0
        for i, cs in enumerate(column_starts):
            if left >= cs - threshold:
                col_idx = i
        columns[col_idx].append(r)
    return columns


def _group_with_rules(regions, page_width):
    columns = _assign_columns(regions, page_width)
    articles = []
    next_id = 1
    for col_idx in sorted(columns.keys()):
        col_regs = sorted(columns[col_idx], key=lambda r: r["bbox"][1])
        current = None
        for reg in col_regs:
            rtype = reg["type"]
            dateline = reg.get("dateline")
            starts = (rtype in ("title", "header")) or (dateline is not None)
            if starts:
                if current is not None and current["member_regions"]:
                    articles.append(current)
                current = {
                    "article_id": next_id,
                    "title_region": reg["id"] if rtype in ("title", "header") else None,
                    "member_regions": [reg["id"]],
                    "bbox": list(reg["bbox"]),
                    "dateline": dateline,
                    "grouped_by": "rules",
                }
                next_id += 1
            else:
                if current is None:
                    current = {
                        "article_id": next_id, "title_region": None,
                        "member_regions": [reg["id"]], "bbox": list(reg["bbox"]),
                        "dateline": None, "grouped_by": "rules",
                    }
                    next_id += 1
                else:
                    current["member_regions"].append(reg["id"])
                    current["bbox"] = _bbox_union(current["bbox"], reg["bbox"])
        if current is not None and current["member_regions"]:
            articles.append(current)
    return articles


def group_into_articles(page_png_path, regions, page_width, page_img=None):
    """Claude first; rules fallback. Returns (articles, enriched_regions)."""
    topband_ocr = _make_topband_ocr(page_img) if page_img else None
    enriched = _enrich_regions_with_text(regions, topband_ocr)

    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if api_key:
        try:
            articles = _group_with_claude(page_png_path, enriched, api_key)
            if articles:
                return articles, enriched
        except Exception as e:
            print(f"[claude grouping failed, falling back to rules]: {e}")

    articles = _group_with_rules(enriched, page_width)
    return articles, enriched


# ---------------------------------------------------------------------------
# Stage 4 — crop each article
# ---------------------------------------------------------------------------
def crop_articles(page_png_path, 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:
        l, t, r, b = art["bbox"]
        margin = 6
        l = max(0, l - margin); t = max(0, t - margin)
        r = min(img.width, r + margin); b = min(img.height, b + margin)

        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))
        results.append({"name": name, "dir": str(art_dir), **info})
    return results


# ---------------------------------------------------------------------------
# Stage 5 — OCR each cropped article (Telugu)
# ---------------------------------------------------------------------------
def _ocr_with_paddleocr(img_path):
    """Use PaddleOCR for better reading-order-aware Telugu OCR."""
    try:
        import paddleocr
        ocr = paddleocr.PaddleOCR(lang='te')
        result = ocr.predict(str(img_path))

        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


def _ocr_with_tesseract(img_path):
    """Fallback: use Tesseract with column-aware PSM for Telugu OCR."""
    try:
        import pytesseract
        # PSM 1 = Automatic page segmentation with OSD (better for multi-column)
        # PSM 3 = Fully automatic (default, often jumbles columns)
        config = '--psm 1 --oem 3'
        text = pytesseract.image_to_string(Image.open(img_path), lang="tel", config=config)
        if not text.strip():
            # Fallback to PSM 4 (single column) if PSM 1 gives nothing
            config = '--psm 4 --oem 3'
            text = pytesseract.image_to_string(Image.open(img_path), lang="tel", config=config)
        return text
    except Exception as e:
        return f"[OCR error: {e}]"


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

        # Try PaddleOCR first (better reading order), fall back to Tesseract
        text = _ocr_with_paddleocr(img_path)
        if not text or len(text.strip()) < 10:
            text = _ocr_with_tesseract(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())
            first_line = next((ln.strip() for ln in text.splitlines() if ln.strip()), "")
            info["headline_preview"] = first_line[:120]
            info_path.write_text(json.dumps(info, indent=2, ensure_ascii=False))


# ---------------------------------------------------------------------------
# Political Filter — crop first, then filter with Claude
# ---------------------------------------------------------------------------
def _filter_articles_with_claude(page_img_path, article_ids_with_locations, page_num, job_dir=None):
    """Send page image + article IDs to Claude, get back only politically significant IDs."""
    import base64
    from io import BytesIO
    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        return None

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

        # Resize page image for API
        img = Image.open(page_img_path)
        max_w = 1500
        if img.width > max_w:
            ratio = max_w / img.width
            img = img.resize((max_w, int(img.height * ratio)), Image.LANCZOS)

        buf = BytesIO()
        img.convert("RGB").save(buf, format="JPEG", quality=80)
        img_data = base64.standard_b64encode(buf.getvalue()).decode("utf-8")

        # Build article list text
        article_list = "\n".join(
            f"- {aid}: located at {loc}" for aid, loc in article_ids_with_locations
        )

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

I have already cropped the following articles from this newspaper page:

{article_list}

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 or ANNOUNCEMENT (not something bad that happened) → answer is NO.
- 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:
  - "id": the article ID
  - "headline_english": English translation
  - "q1": "YES" or "NO"
  - "q1_reason": why (1 sentence)
  - "q2": "YES" or "NO" (skip if q1 is NO)
  - "q2_reason": why (1 sentence, 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": your answer to Question 1 (1 sentence)
  - "q2_answer": your answer to Question 2 (1 sentence)
  - "attack_angle": what opposition should say in a press conference (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-opus-4-8", max_tokens=4096,
            messages=[{"role": "user", "content": [
                {"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": img_data}},
                {"type": "text", "text": 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
        from pathlib import Path
        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):
    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)

    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)
        )
        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)
        )

        # --- 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 ALL articles with precise PaddleOCR boundaries
        # _set_status(job_id, f"Page {n}: cropping {len(articles)} articles...")
        # records = crop_articles(page_img_path, articles, articles_root, n)
        # print(f"  Page {n}: cropped {len(records)} articles")
        #
        # # Build article ID + location + headline list for Claude
        # article_ids_with_locations = []
        # for rec in records:
        #     art_dir = Path(rec["dir"])
        #     art_name = art_dir.name
        #     info_path = art_dir / "info.json"
        #     if info_path.exists():
        #         info = json.loads(info_path.read_text())
        #         bbox = info.get("bbox", [0, 0, 0, 0])
        #         hint = info.get("headline_hint", "")
        #         pw, ph = page_img.size
        #         loc = f"headline: \"{hint}\", position: x:{bbox[0]}-{bbox[2]}, y:{bbox[1]}-{bbox[3]}"
        #     else:
        #         loc = "unknown"
        #     article_ids_with_locations.append((art_name, loc))
        #
        # # Step 5: Political filter — send page image + article IDs to Claude
        # if article_ids_with_locations:
        #     _set_status(job_id, f"Page {n}: filtering {len(records)} articles for political significance...")
        #     selected = _filter_articles_with_claude(page_img_path, article_ids_with_locations, n, job_dir=str(job_dir))
        #
        #     if selected:
        #         selected_ids = {s["id"] for s in selected}
        #         print(f"  Page {n}: {len(selected)} politically significant articles: {selected_ids}")
        #
        #         # Step 6: 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 selected_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())
        #                     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))
        #                 all_records.append(rec)
        #                 all_selected.append(sel_info)
        #             else:
        #                 # Delete non-political article
        #                 import shutil
        #                 shutil.rmtree(str(art_dir), ignore_errors=True)
        #     else:
        #         # No political articles on this page — delete all cropped articles
        #         print(f"  Page {n}: no political articles found, removing all {len(records)} articles")
        #         import shutil
        #         for rec in records:
        #             shutil.rmtree(str(Path(rec["dir"])), ignore_errors=True)
        # else:
        #     all_records.extend(records)

        print(f"  Page {n}: Steps 3-6 skipped (commented out). Check regions JSON.")

    # --- 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())
            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
    # all_images_dir = job_dir / "all_political_images"
    # all_images_dir.mkdir(parents=True, exist_ok=True)
    # import shutil
    # 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}")

    # # Create a single PDF with all political article images
    # if all_records:
    #     pdf_images = []
    #     for rec in sorted(all_records, key=lambda r: r.get("article_id", 0)):
    #         art_dir = Path(rec["dir"])
    #         img_path = art_dir / "article.png"
    #         if img_path.exists():
    #             pdf_images.append(Image.open(img_path).convert("RGB"))
    #     if pdf_images:
    #         pdf_path_out = job_dir / "political_articles.pdf"
    #         pdf_images[0].save(str(pdf_path_out), save_all=True, append_images=pdf_images[1:])
    #         print(f"  Created PDF with {len(pdf_images)} articles: {pdf_path_out}")

    _set_status(job_id, f"Done! Steps 1-2 complete. {len(regions)} regions detected.")
    return {"pages": len(pages), "articles": 0}
