"""
OCR Processing for Scanned Documents

Uses Tesseract via pytesseract for optical character recognition.
Includes pre-processing for improved accuracy and post-processing
for technical term correction.
"""

import logging
import re
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path

logger = logging.getLogger(__name__)


class OCREngine(Enum):
    """Available OCR engines."""

    TESSERACT = "tesseract"
    # Future: EASYOCR = "easyocr"
    # Future: PADDLEOCR = "paddleocr"


@dataclass
class OCRConfig:
    """Configuration for OCR processing."""

    engine: OCREngine = OCREngine.TESSERACT
    language: str = "eng"
    page_segmentation_mode: int = 3  # Fully automatic page segmentation
    oem: int = 3  # LSTM OCR Engine Mode
    dpi: int = 300
    # Pre-processing options
    deskew: bool = True
    denoise: bool = True
    enhance_contrast: bool = True
    binarize: bool = True
    # Post-processing options
    spell_correct: bool = True
    preserve_technical_terms: bool = True


@dataclass
class OCRResult:
    """Result of OCR processing."""

    text: str
    confidence: float  # 0.0 to 1.0
    page_number: int | None = None
    word_confidences: list[tuple[str, float]] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)

    @property
    def is_high_quality(self) -> bool:
        """Check if OCR result is high quality (>85% confidence)."""
        return self.confidence >= 0.85

    @property
    def is_acceptable(self) -> bool:
        """Check if OCR result is acceptable (>70% confidence)."""
        return self.confidence >= 0.70


@dataclass
class OCRPageResult:
    """OCR result for a single page."""

    page_number: int
    text: str
    confidence: float
    word_count: int


@dataclass
class OCRDocumentResult:
    """OCR result for an entire document."""

    file_path: Path
    total_pages: int
    pages: list[OCRPageResult]
    overall_confidence: float
    full_text: str
    warnings: list[str] = field(default_factory=list)

    @property
    def low_confidence_pages(self) -> list[int]:
        """Get page numbers with low confidence."""
        return [p.page_number for p in self.pages if p.confidence < 0.70]


# Technical terms that should not be spell-corrected
NCS_TECHNICAL_TERMS = {
    # Isotopes
    "U-235", "U-238", "Pu-239", "Pu-240", "Pu-241", "Cf-252",
    "Am-241", "Np-237", "Cm-244", "U-233",
    # Units and values
    "keff", "k-eff", "k_eff", "kinf", "k-inf",
    "MeV", "keV", "eV", "barn", "barns", "mb",
    "wt%", "at%", "g/cm3", "kg/m3",
    # Codes
    "MCNP", "SCALE", "KENO", "ENDF", "JEFF", "JENDL",
    "TSUNAMI", "CSAS", "MAVRIC", "ORIGEN",
    # Organizations
    "ANSI", "ANS", "NRC", "DOE", "NNSA", "ORNL", "LANL", "SNL",
    # Technical terms
    "subcritical", "supercritical", "fissile", "fissionable",
    "moderator", "reflector", "absorber", "neutron",
    "enrichment", "burnup", "reactivity", "eigenvalue",
    "NUREG", "ICSBEP", "OECD/NEA",
}


class OCRProcessor:
    """
    OCR processor for scanned documents.

    Features:
    - Pre-processing: deskew, denoise, contrast enhancement
    - Multiple OCR engine support (currently Tesseract)
    - Post-processing: spell correction with technical term preservation
    - Confidence scoring

    Example:
        processor = OCRProcessor()
        result = processor.process_pdf("scanned_document.pdf")
        print(f"Confidence: {result.overall_confidence:.1%}")
        print(result.full_text)
    """

    def __init__(self, config: OCRConfig | None = None):
        """
        Initialize OCR processor.

        Args:
            config: OCR configuration (uses defaults if not provided)
        """
        self.config = config or OCRConfig()
        self._check_dependencies()

    def _check_dependencies(self):
        """Check that required dependencies are available."""
        if self.config.engine == OCREngine.TESSERACT:
            try:
                import pytesseract
                # Check Tesseract is installed
                pytesseract.get_tesseract_version()
            except ImportError:
                raise ImportError("pytesseract not installed. Run: pip install pytesseract") from None
            except Exception as e:
                raise RuntimeError(f"Tesseract not found. Install Tesseract OCR: {e}") from e

    def process_image(
        self,
        image_path: str | Path,
        page_number: int | None = None,
    ) -> OCRResult:
        """
        Process a single image with OCR.

        Args:
            image_path: Path to image file
            page_number: Optional page number for tracking

        Returns:
            OCRResult with extracted text and confidence
        """
        import pytesseract
        from PIL import Image

        image_path = Path(image_path)

        # Load image
        image = Image.open(image_path)

        # Pre-process image
        if self.config.deskew or self.config.denoise or self.config.enhance_contrast:
            image = self._preprocess_image(image)

        # Configure Tesseract
        custom_config = (
            f"--psm {self.config.page_segmentation_mode} "
            f"--oem {self.config.oem}"
        )

        # Run OCR
        try:
            # Get detailed data for confidence calculation
            data = pytesseract.image_to_data(
                image,
                lang=self.config.language,
                config=custom_config,
                output_type=pytesseract.Output.DICT,
            )

            # Calculate word confidences
            word_confidences = []
            words = []
            confidences = []

            for i, word in enumerate(data["text"]):
                conf = data["conf"][i]
                if word.strip() and conf != -1:
                    words.append(word)
                    confidences.append(conf / 100.0)
                    word_confidences.append((word, conf / 100.0))

            # Calculate overall confidence
            overall_confidence = sum(confidences) / len(confidences) if confidences else 0.0

            # Get full text
            text = pytesseract.image_to_string(
                image,
                lang=self.config.language,
                config=custom_config,
            )

            # Post-process text
            if self.config.spell_correct:
                text = self._postprocess_text(text)

            warnings = []
            if overall_confidence < 0.70:
                warnings.append(f"Low OCR confidence: {overall_confidence:.1%}")

            return OCRResult(
                text=text,
                confidence=overall_confidence,
                page_number=page_number,
                word_confidences=word_confidences,
                warnings=warnings,
            )

        except Exception as e:
            logger.error(f"OCR failed for {image_path}: {e}")
            return OCRResult(
                text="",
                confidence=0.0,
                page_number=page_number,
                warnings=[f"OCR failed: {str(e)}"],
            )

    def process_pdf(
        self,
        pdf_path: str | Path,
        pages: list[int] | None = None,
    ) -> OCRDocumentResult:
        """
        Process a PDF document with OCR.

        Args:
            pdf_path: Path to PDF file
            pages: Specific pages to process (1-indexed). None = all pages.

        Returns:
            OCRDocumentResult with text from all pages
        """
        try:
            import fitz  # pymupdf  # noqa: F401
        except ImportError:
            raise ImportError("pymupdf not installed. Run: pip install pymupdf") from None

        import io

        from PIL import Image

        pdf_path = Path(pdf_path)
        doc = fitz.open(pdf_path)

        page_results = []
        all_text = []
        all_confidences = []
        warnings = []

        try:
            # Determine which pages to process
            if pages is None:
                page_indices = list(range(len(doc)))
            else:
                page_indices = [p - 1 for p in pages if 0 < p <= len(doc)]

            for page_idx in page_indices:
                page = doc[page_idx]
                page_num = page_idx + 1

                # Convert page to image
                mat = fitz.Matrix(self.config.dpi / 72, self.config.dpi / 72)
                pix = page.get_pixmap(matrix=mat)
                img_data = pix.tobytes("png")

                # Convert to PIL Image
                image = Image.open(io.BytesIO(img_data))

                # Pre-process
                if self.config.deskew or self.config.denoise or self.config.enhance_contrast:
                    image = self._preprocess_image(image)

                # Run OCR on this page
                result = self._ocr_pil_image(image, page_num)

                page_results.append(OCRPageResult(
                    page_number=page_num,
                    text=result.text,
                    confidence=result.confidence,
                    word_count=len(result.text.split()),
                ))

                all_text.append(f"\n--- Page {page_num} ---\n")
                all_text.append(result.text)
                all_confidences.append(result.confidence)

                if result.warnings:
                    warnings.extend([f"Page {page_num}: {w}" for w in result.warnings])

            overall_confidence = (
                sum(all_confidences) / len(all_confidences)
                if all_confidences else 0.0
            )

            return OCRDocumentResult(
                file_path=pdf_path,
                total_pages=len(doc),
                pages=page_results,
                overall_confidence=overall_confidence,
                full_text="\n".join(all_text),
                warnings=warnings,
            )

        finally:
            doc.close()

    def _ocr_pil_image(self, image, page_number: int) -> OCRResult:
        """Run OCR on a PIL Image object."""
        import pytesseract

        custom_config = (
            f"--psm {self.config.page_segmentation_mode} "
            f"--oem {self.config.oem}"
        )

        try:
            # Get detailed data for confidence
            data = pytesseract.image_to_data(
                image,
                lang=self.config.language,
                config=custom_config,
                output_type=pytesseract.Output.DICT,
            )

            confidences = [
                c / 100.0 for c in data["conf"]
                if c != -1 and data["text"][data["conf"].index(c)].strip()
            ]

            overall_confidence = sum(confidences) / len(confidences) if confidences else 0.0

            text = pytesseract.image_to_string(
                image,
                lang=self.config.language,
                config=custom_config,
            )

            if self.config.spell_correct:
                text = self._postprocess_text(text)

            warnings = []
            if overall_confidence < 0.70:
                warnings.append(f"Low confidence: {overall_confidence:.1%}")

            return OCRResult(
                text=text,
                confidence=overall_confidence,
                page_number=page_number,
                warnings=warnings,
            )

        except Exception as e:
            logger.error(f"OCR failed for page {page_number}: {e}")
            return OCRResult(
                text="",
                confidence=0.0,
                page_number=page_number,
                warnings=[f"OCR failed: {str(e)}"],
            )

    def _preprocess_image(self, image):
        """Apply pre-processing to improve OCR quality."""
        from PIL import ImageEnhance, ImageFilter

        # Convert to grayscale if needed
        if image.mode != "L":
            image = image.convert("L")

        # Enhance contrast
        if self.config.enhance_contrast:
            enhancer = ImageEnhance.Contrast(image)
            image = enhancer.enhance(1.5)

        # Denoise
        if self.config.denoise:
            image = image.filter(ImageFilter.MedianFilter(size=3))

        # Binarize (convert to black and white)
        if self.config.binarize:
            threshold = 128
            image = image.point(lambda x: 255 if x > threshold else 0, "1")
            image = image.convert("L")

        # Deskew would require additional libraries like deskew or opencv
        # Skipping for now as it adds complexity

        return image

    def _postprocess_text(self, text: str) -> str:
        """Post-process OCR text to correct common errors."""
        if not text:
            return text

        # Protect technical terms
        protected_terms = {}
        for i, term in enumerate(NCS_TECHNICAL_TERMS):
            pattern = re.compile(re.escape(term), re.IGNORECASE)
            if pattern.search(text):
                placeholder = f"__PROTECTED_{i}__"
                protected_terms[placeholder] = term
                text = pattern.sub(placeholder, text)

        # Common OCR substitution errors
        corrections = {
            r"\bl\b": "1",  # Isolated 'l' often means '1'
            r"\bO\b": "0",  # Isolated 'O' often means '0'
            r"rn": "m",  # 'rn' often misread for 'm'
            r"vv": "w",  # 'vv' often misread for 'w'
            r"\s+": " ",  # Normalize whitespace
        }

        # Apply corrections carefully
        for _pattern, _replacement in corrections.items():
            # Only apply if not in a word context
            pass  # Skip aggressive corrections to preserve content

        # Normalize whitespace
        text = re.sub(r"\s+", " ", text)
        text = re.sub(r"\n{3,}", "\n\n", text)

        # Restore protected technical terms
        for placeholder, term in protected_terms.items():
            text = text.replace(placeholder, term)

        return text.strip()


def process_scanned_pdf(
    pdf_path: str | Path,
    config: OCRConfig | None = None,
) -> OCRDocumentResult:
    """
    Convenience function to process a scanned PDF.

    Args:
        pdf_path: Path to PDF file
        config: Optional OCR configuration

    Returns:
        OCRDocumentResult with extracted text
    """
    processor = OCRProcessor(config)
    return processor.process_pdf(pdf_path)
