"""
OCR Extraction Module for Baymax
================================
Extracts text from scanned PDFs and images using Tesseract OCR.
Automatically detects if a document needs OCR or has embedded text.
"""

import os
import re
import tempfile
from typing import Optional, Tuple, List
from pathlib import Path

# ============================================================================
# DEPENDENCY CHECKS
# ============================================================================

# Check for pytesseract (OCR engine)
try:
    import pytesseract
    from PIL import Image

    # Try to find Tesseract installation
    tesseract_paths = [
        r"C:\Program Files\Tesseract-OCR\tesseract.exe",
        r"C:\Program Files (x86)\Tesseract-OCR\tesseract.exe",
        os.path.expanduser(r"~\AppData\Local\Programs\Tesseract-OCR\tesseract.exe"),
        os.path.expanduser(r"~\AppData\Local\Tesseract-OCR\tesseract.exe"),
        r"C:\Tesseract-OCR\tesseract.exe",
        "/usr/bin/tesseract",  # Linux
        "/usr/local/bin/tesseract",  # macOS
    ]

    TESSERACT_AVAILABLE = False
    for path in tesseract_paths:
        if os.path.exists(path):
            pytesseract.pytesseract.tesseract_cmd = path
            try:
                pytesseract.get_tesseract_version()
                TESSERACT_AVAILABLE = True
                TESSERACT_PATH = path
                break
            except Exception:
                continue

    # Try system PATH as fallback
    if not TESSERACT_AVAILABLE:
        try:
            pytesseract.get_tesseract_version()
            TESSERACT_AVAILABLE = True
            TESSERACT_PATH = "tesseract (in PATH)"
        except Exception:
            pass

except ImportError:
    TESSERACT_AVAILABLE = False
    TESSERACT_PATH = None
    pytesseract = None
    Image = None

# Check for pdf2image (PDF to image conversion)
try:
    from pdf2image import convert_from_path
    from pdf2image.exceptions import PDFInfoNotInstalledError

    # Check for poppler (required by pdf2image on Windows)
    poppler_paths = [
        os.path.join(os.path.dirname(os.path.abspath(__file__)), "poppler-25.12.0", "Library", "bin"),
        os.path.join(os.path.dirname(os.path.abspath(__file__)), "poppler-24.08.0", "Library", "bin"),
        os.path.join(os.path.dirname(os.path.abspath(__file__)), "poppler", "Library", "bin"),
        r"C:\Program Files\poppler-25.12.0\Library\bin",
        r"C:\Program Files\poppler-24.08.0\Library\bin",
        r"C:\Program Files\poppler\Library\bin",
        r"C:\poppler\Library\bin",
    ]

    POPPLER_PATH = None
    for path in poppler_paths:
        if os.path.exists(path):
            POPPLER_PATH = path
            break

    PDF2IMAGE_AVAILABLE = True
except ImportError:
    PDF2IMAGE_AVAILABLE = False
    POPPLER_PATH = None
    convert_from_path = None

# Check for pypdf (text extraction from PDFs)
try:
    from pypdf import PdfReader
    PYPDF_AVAILABLE = True
except ImportError:
    PYPDF_AVAILABLE = False
    PdfReader = None


# ============================================================================
# OCR CONFIGURATION
# ============================================================================

# Tesseract configuration for legal documents
# PSM 3 = Fully automatic page segmentation (default)
# PSM 6 = Assume uniform block of text
TESSERACT_CONFIG = '--psm 3 --oem 3'

# Language (eng = English, can add more like 'eng+fra' for multilingual)
TESSERACT_LANG = 'eng'


# ============================================================================
# UTILITY FUNCTIONS
# ============================================================================

def get_ocr_status() -> dict:
    """Get the status of OCR dependencies."""
    return {
        "tesseract_available": TESSERACT_AVAILABLE,
        "tesseract_path": TESSERACT_PATH if TESSERACT_AVAILABLE else None,
        "pdf2image_available": PDF2IMAGE_AVAILABLE,
        "poppler_path": POPPLER_PATH,
        "pypdf_available": PYPDF_AVAILABLE,
        "fully_operational": TESSERACT_AVAILABLE and (PDF2IMAGE_AVAILABLE or PYPDF_AVAILABLE),
    }


def is_image_file(file_path: str) -> bool:
    """Check if file is an image."""
    image_extensions = {'.png', '.jpg', '.jpeg', '.tiff', '.tif', '.bmp', '.gif', '.webp'}
    return Path(file_path).suffix.lower() in image_extensions


def is_pdf_file(file_path: str) -> bool:
    """Check if file is a PDF."""
    return Path(file_path).suffix.lower() == '.pdf'


def clean_ocr_text(text: str) -> str:
    """Clean up OCR output text."""
    # Remove excessive whitespace
    text = re.sub(r'\n{3,}', '\n\n', text)
    text = re.sub(r' {2,}', ' ', text)

    # Fix common OCR errors in legal text
    replacements = [
        (r'\bll\b', 'll'),  # Common misread
        (r'(?<=[a-z])1(?=[a-z])', 'l'),  # 1 instead of l in words
        (r'(?<=[a-z])0(?=[a-z])', 'o'),  # 0 instead of o in words
        (r'\bU\.S\. C\.', 'U.S.C.'),  # Fix statute citations
        (r'\bF\. ?2d\b', 'F.2d'),
        (r'\bF\. ?3d\b', 'F.3d'),
        (r'\bS\. ?Ct\.', 'S.Ct.'),
    ]

    for pattern, replacement in replacements:
        text = re.sub(pattern, replacement, text)

    return text.strip()


# ============================================================================
# PDF TEXT EXTRACTION (non-OCR)
# ============================================================================

def extract_text_from_pdf(file_path: str) -> Tuple[str, bool]:
    """
    Try to extract embedded text from PDF (no OCR).

    Returns:
        Tuple of (extracted_text, needs_ocr)
        If the PDF has embedded text, needs_ocr=False
        If the PDF is scanned/image-based, needs_ocr=True
    """
    if not PYPDF_AVAILABLE:
        return "", True

    try:
        reader = PdfReader(file_path)
        all_text = []

        for page in reader.pages:
            text = page.extract_text() or ""
            all_text.append(text)

        full_text = "\n\n".join(all_text)

        # Heuristic: if we got very little text, it's probably scanned
        # A typical legal page has 300-500 words
        word_count = len(full_text.split())
        page_count = len(reader.pages)

        if page_count > 0 and word_count / page_count < 50:
            # Less than 50 words per page suggests scanned document
            return full_text, True

        return full_text, False

    except Exception as e:
        return f"Error reading PDF: {str(e)}", True


# ============================================================================
# OCR EXTRACTION
# ============================================================================

def ocr_image(image_path: str) -> str:
    """
    Extract text from an image file using OCR.

    Args:
        image_path: Path to image file

    Returns:
        Extracted text
    """
    if not TESSERACT_AVAILABLE:
        return "ERROR: Tesseract OCR is not installed. Please install from: https://github.com/UB-Mannheim/tesseract/wiki"

    try:
        image = Image.open(image_path)

        # Convert to RGB if necessary (handles RGBA, grayscale, etc.)
        if image.mode != 'RGB':
            image = image.convert('RGB')

        # Run OCR
        text = pytesseract.image_to_string(
            image,
            lang=TESSERACT_LANG,
            config=TESSERACT_CONFIG
        )

        return clean_ocr_text(text)

    except Exception as e:
        return f"ERROR: Failed to OCR image: {str(e)}"


def ocr_pdf(file_path: str, max_pages: int = 50) -> str:
    """
    Extract text from a scanned PDF using OCR.

    Args:
        file_path: Path to PDF file
        max_pages: Maximum pages to process (to prevent timeout on huge docs)

    Returns:
        Extracted text
    """
    if not TESSERACT_AVAILABLE:
        return "ERROR: Tesseract OCR is not installed. Please install from: https://github.com/UB-Mannheim/tesseract/wiki"

    if not PDF2IMAGE_AVAILABLE:
        return "ERROR: pdf2image is not installed. Run: pip install pdf2image"

    try:
        # Convert PDF to images
        convert_kwargs = {"first_page": 1, "last_page": max_pages}

        if POPPLER_PATH:
            convert_kwargs["poppler_path"] = POPPLER_PATH

        images = convert_from_path(file_path, **convert_kwargs)

        all_text = []
        for i, image in enumerate(images, 1):
            # Convert to RGB if needed
            if image.mode != 'RGB':
                image = image.convert('RGB')

            # OCR this page
            page_text = pytesseract.image_to_string(
                image,
                lang=TESSERACT_LANG,
                config=TESSERACT_CONFIG
            )

            all_text.append(f"--- Page {i} ---\n{page_text}")

        return clean_ocr_text("\n\n".join(all_text))

    except PDFInfoNotInstalledError:
        return "ERROR: Poppler is not installed. On Windows, download from: https://github.com/oschwartz10612/poppler-windows/releases"
    except Exception as e:
        return f"ERROR: Failed to OCR PDF: {str(e)}"


# ============================================================================
# MAIN EXTRACTION FUNCTION
# ============================================================================

def extract_text_with_ocr(file_path: str, force_ocr: bool = False) -> dict:
    """
    Extract text from a document, using OCR if necessary.

    This is the main entry point. It automatically:
    1. Detects file type (PDF vs image)
    2. For PDFs: tries embedded text first, falls back to OCR
    3. For images: uses OCR directly

    Args:
        file_path: Path to the document
        force_ocr: If True, skip embedded text extraction and go straight to OCR

    Returns:
        dict with:
        - text: The extracted text
        - method: "embedded" or "ocr"
        - pages: Number of pages processed (for PDFs)
        - success: True if extraction worked
        - error: Error message if any
    """
    result = {
        "text": "",
        "method": "unknown",
        "pages": 0,
        "success": False,
        "error": None,
        "file_path": file_path,
    }

    # Check file exists
    if not os.path.exists(file_path):
        result["error"] = f"File not found: {file_path}"
        return result

    # Handle image files
    if is_image_file(file_path):
        result["method"] = "ocr"
        result["pages"] = 1

        text = ocr_image(file_path)
        if text.startswith("ERROR:"):
            result["error"] = text
        else:
            result["text"] = text
            result["success"] = True

        return result

    # Handle PDF files
    if is_pdf_file(file_path):
        # Try embedded text first (unless force_ocr)
        if not force_ocr:
            embedded_text, needs_ocr = extract_text_from_pdf(file_path)

            if not needs_ocr and embedded_text:
                result["method"] = "embedded"
                result["text"] = embedded_text
                result["success"] = True
                # Count pages
                try:
                    reader = PdfReader(file_path)
                    result["pages"] = len(reader.pages)
                except:
                    pass
                return result

        # Need OCR
        result["method"] = "ocr"
        text = ocr_pdf(file_path)

        if text.startswith("ERROR:"):
            result["error"] = text
        else:
            result["text"] = text
            result["success"] = True
            # Count pages from OCR output
            result["pages"] = text.count("--- Page ")

        return result

    # Unsupported file type
    result["error"] = f"Unsupported file type: {Path(file_path).suffix}"
    return result


# ============================================================================
# MCP TOOL FUNCTION
# ============================================================================

def ocr_document_for_mcp(file_path: str, force_ocr: bool = False) -> str:
    """
    MCP-friendly wrapper for document OCR.

    Args:
        file_path: Path to the document (PDF or image)
        force_ocr: If True, always use OCR even for PDFs with embedded text

    Returns:
        Formatted string with extraction results
    """
    # First check OCR status
    status = get_ocr_status()

    if not status["fully_operational"]:
        output = ["OCR DEPENDENCY STATUS", "=" * 70]
        output.append(f"Tesseract OCR: {'INSTALLED' if status['tesseract_available'] else 'NOT FOUND'}")
        if status['tesseract_path']:
            output.append(f"  Path: {status['tesseract_path']}")
        output.append(f"PDF Support: {'AVAILABLE' if status['pdf2image_available'] else 'NOT AVAILABLE'}")
        if status['poppler_path']:
            output.append(f"  Poppler: {status['poppler_path']}")

        output.append("\nINSTALLATION INSTRUCTIONS:")
        if not status['tesseract_available']:
            output.append("  Tesseract: https://github.com/UB-Mannheim/tesseract/wiki")
        if not status['pdf2image_available']:
            output.append("  pdf2image: pip install pdf2image")
            output.append("  Poppler (Windows): https://github.com/oschwartz10612/poppler-windows/releases")

        return "\n".join(output)

    # Run extraction
    result = extract_text_with_ocr(file_path, force_ocr)

    output = ["DOCUMENT TEXT EXTRACTION", "=" * 70]
    output.append(f"File: {result['file_path']}")
    output.append(f"Method: {result['method'].upper()}")
    output.append(f"Pages: {result['pages']}")
    output.append(f"Success: {'YES' if result['success'] else 'NO'}")

    if result['error']:
        output.append(f"\nERROR: {result['error']}")

    if result['success']:
        output.append("\n" + "-" * 70)
        output.append("EXTRACTED TEXT:")
        output.append("-" * 70)
        output.append(result['text'])

    return "\n".join(output)


# ============================================================================
# QUICK TEST
# ============================================================================

if __name__ == "__main__":
    print("OCR Extraction Module - Status Check")
    print("=" * 50)

    status = get_ocr_status()
    for key, value in status.items():
        print(f"{key}: {value}")

    if status["fully_operational"]:
        print("\nOCR system is fully operational!")
    else:
        print("\nSome components are missing. See above for details.")
