"""
Read Telugu text from cropped article images using Claude Vision API.
Run after a job completes to get accurate text extraction.

Usage:
  python claude_ocr.py                          # processes latest job
  python claude_ocr.py 20260526_002006          # processes specific job
"""
import os
import sys
import json
import base64
import anthropic
from pathlib import Path

OUTPUT_DIR = Path(__file__).parent / "output"


def read_article_image(client, img_path, model="claude-sonnet-4-6"):
    """Send article image to Claude and get Telugu text back."""
    with open(img_path, "rb") as f:
        img_data = base64.standard_b64encode(f.read()).decode("utf-8")

    response = client.messages.create(
        model=model,
        max_tokens=4096,
        messages=[{
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "source": {
                        "type": "base64",
                        "media_type": "image/png",
                        "data": img_data,
                    },
                },
                {
                    "type": "text",
                    "text": (
                        "Read all Telugu text from this newspaper article image. "
                        "Output ONLY the Telugu text in correct reading order "
                        "(top to bottom, reading each column fully before moving to the next). "
                        "Include headlines first, then subheadlines, then body text. "
                        "Do not translate. Do not add any explanation. Just the Telugu text."
                    ),
                },
            ],
        }],
    )
    return response.content[0].text


def process_job(job_id=None):
    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        print("Error: Set ANTHROPIC_API_KEY environment variable")
        sys.exit(1)

    # Find job directory
    if job_id:
        job_dir = OUTPUT_DIR / job_id
    else:
        # Use latest job
        jobs = sorted(OUTPUT_DIR.iterdir(), reverse=True)
        if not jobs:
            print("No jobs found")
            sys.exit(1)
        job_dir = jobs[0]
        job_id = job_dir.name

    articles_dir = job_dir / "articles"
    if not articles_dir.exists():
        print(f"No articles found in {job_dir}")
        sys.exit(1)

    client = anthropic.Anthropic(api_key=api_key)
    article_dirs = sorted(articles_dir.iterdir())

    print(f"Job: {job_id}")
    print(f"Articles: {len(article_dirs)}")
    print("-" * 50)

    all_text = []

    for art_dir in article_dirs:
        img_path = art_dir / "article.png"
        if not img_path.exists():
            continue

        name = art_dir.name
        print(f"Reading {name}...", end=" ", flush=True)

        try:
            text = read_article_image(client, img_path)
            # Save to article.txt (overwrites old OCR output)
            (art_dir / "article.txt").write_text(text, encoding="utf-8")
            print(f"OK ({len(text)} chars)")

            all_text.append(f"--- {name} ---\n{text}\n")

            # Update info.json headline preview
            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["ocr_method"] = "claude_vision"
                info_path.write_text(json.dumps(info, indent=2, ensure_ascii=False))

        except Exception as e:
            print(f"FAILED: {e}")
            all_text.append(f"--- {name} ---\n[Error: {e}]\n")

    # Save combined file
    combined_path = job_dir / "all_articles_claude.txt"
    combined_path.write_text("\n".join(all_text), encoding="utf-8")
    print("-" * 50)
    print(f"Done! Combined text saved to: {combined_path}")


if __name__ == "__main__":
    job_id = sys.argv[1] if len(sys.argv) > 1 else None
    process_job(job_id)