"""
Extract text from a screenshot using Tesseract OCR.

Usage:
    python ocr.py --input screenshot.png
    python ocr.py --input screenshot.png --region 100,200,400,50
    python ocr.py --input screenshot.png --output results.json --lang eng
"""
import argparse
import json
import os
import sys

try:
    from PIL import Image
except ImportError:
    print("Error: Pillow is not installed. Run the setup script first.", file=sys.stderr)
    sys.exit(1)

try:
    import pytesseract
except ImportError:
    print("Error: pytesseract is not installed. Run the setup script first.", file=sys.stderr)
    sys.exit(1)


def extract_text(image_path, region=None, lang="eng"):
    """
    Run OCR on an image and return structured results.

    Args:
        image_path: Path to the PNG image
        region: Optional (x, y, w, h) tuple to crop before OCR
        lang: Tesseract language code

    Returns:
        dict with 'full_text' and 'blocks' (bounding boxes + text)
    """
    img = Image.open(image_path)

    if region:
        x, y, w, h = region
        img = img.crop((x, y, x + w, y + h))

    # Get full text
    full_text = pytesseract.image_to_string(img, lang=lang).strip()

    # Get detailed bounding box data
    data = pytesseract.image_to_data(img, lang=lang, output_type=pytesseract.Output.DICT)

    blocks = []
    n = len(data["text"])
    for i in range(n):
        text = data["text"][i].strip()
        conf = float(data["conf"][i])
        if text and conf > 0:
            block = {
                "text": text,
                "x": data["left"][i],
                "y": data["top"][i],
                "width": data["width"][i],
                "height": data["height"][i],
                "confidence": round(conf, 1),
            }
            # If we cropped, offset coordinates back to original image space
            if region:
                block["x"] += region[0]
                block["y"] += region[1]
            blocks.append(block)

    return {"full_text": full_text, "blocks": blocks}


def main():
    parser = argparse.ArgumentParser(description="OCR a screenshot image.")
    parser.add_argument("--input", required=True, help="Input PNG file path")
    parser.add_argument("--output", default=None, help="Output JSON file path (default: <input>_ocr.json)")
    parser.add_argument("--region", default=None, help="Crop region as X,Y,W,H before OCR")
    parser.add_argument("--lang", default="eng", help="Tesseract language code (default: eng)")
    args = parser.parse_args()

    if not os.path.isfile(args.input):
        print(f"Error: Input file not found: {args.input}", file=sys.stderr)
        sys.exit(1)

    region = None
    if args.region:
        try:
            parts = [int(x.strip()) for x in args.region.split(",")]
            if len(parts) != 4:
                raise ValueError()
            region = tuple(parts)
        except ValueError:
            print("Error: --region must be four comma-separated integers: X,Y,W,H", file=sys.stderr)
            sys.exit(1)

    try:
        result = extract_text(args.input, region=region, lang=args.lang)
    except Exception as e:
        print(f"Error running OCR: {e}", file=sys.stderr)
        if "tesseract is not installed" in str(e).lower() or "tesseractnotfounderror" in type(e).__name__.lower():
            print("Tesseract OCR is not installed or not on PATH.", file=sys.stderr)
            print("Install from: https://github.com/UB-Mannheim/tesseract/wiki", file=sys.stderr)
        sys.exit(1)

    # Determine output path
    output_path = args.output
    if not output_path:
        base, _ = os.path.splitext(args.input)
        output_path = f"{base}_ocr.json"

    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(result, f, indent=2)

    # Print full text to stdout for easy piping
    print(result["full_text"])
    print(f"\n--- OCR details saved to: {output_path} ---", file=sys.stderr)


if __name__ == "__main__":
    main()