import pytesseract
from PIL import Image
import io
from typing import Dict, List, Any


class TextExtractionAgent:
    """PDF text extraction using OCR via Tesseract."""

    def __init__(self, tessdata_path: str = None):
        self.tessdata_path = tessdata_path
        if tessdata_path:
            pytesseract.pytesseract.tesseract_cmd = tessdata_path

    def process(self, pdf_page: Image.Image) -> Dict[str, Any]:
        """Extract text from a PDF page image."""
        text = pytesseract.image_to_string(pdf_page)
        return {
            "text": text,
            "confidence": self._get_confidence(pdf_page),
            "blocks": self._extract_blocks(text),
        }

    def _get_confidence(self, image: Image.Image) -> float:
        """Get OCR confidence score."""
        data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)
        confidences = [int(conf) for conf in data["conf"] if conf != "-1"]
        return sum(confidences) / len(confidences) if confidences else 0.0

    def _extract_blocks(self, text: str) -> List[Dict]:
        """Extract text blocks."""
        blocks = []
        for i, line in enumerate(text.split("\n")):
            if line.strip():
                blocks.append({"id": i, "text": line})
        return blocks