#!/usr/bin/env python3
"""
Unit tests for the pydocai extractor module.
"""

import pytest
from unittest.mock import patch, MagicMock
import os

from pydocai.extractor import (
    extract_pdf_text,
    _extract_text_with_pypdfium2,
    _extract_text_with_ocr,
    _has_sufficient_text,
)


@patch("pydocai.extractor._extract_text_with_pypdfium2")
def test_extract_pdf_text_success_with_pypdfium2(
    mock_pypdfium2, temp_pdf_path, temp_txt_path
):
    """Test successful PDF text extraction using pypdfium2."""
    # Mock pypdfium2 to return text with sufficient content (>= 5 lines/page)
    sufficient_text = "--- Page 1 ---\n" + "\n".join(
        [f"Extracted text line {i}" for i in range(10)]
    )
    mock_pypdfium2.return_value = sufficient_text

    result = extract_pdf_text(temp_pdf_path, temp_txt_path)

    assert result is True
    assert os.path.exists(temp_txt_path)

    # Verify the extracted text was saved
    with open(temp_txt_path, "r") as f:
        content = f.read()
        assert "Extracted text line" in content
        assert "--- Page 1 ---" in content

    # Verify pypdfium2 was called
    mock_pypdfium2.assert_called_once_with(temp_pdf_path)


@patch("pydocai.extractor._extract_text_with_ocr")
@patch("pydocai.extractor._extract_text_with_pypdfium2")
def test_extract_pdf_text_fallback_to_ocr(
    mock_pypdfium2, mock_ocr, temp_pdf_path, temp_txt_path
):
    """Test PDF text extraction falls back to OCR when pypdfium2 fails."""
    # Mock pypdfium2 to return None (no text extracted)
    mock_pypdfium2.return_value = None

    # Mock OCR to return text with sufficient content
    ocr_text = "--- Page 1 ---\n" + "\n".join([f"OCR text line {i}" for i in range(10)])
    mock_ocr.return_value = ocr_text

    result = extract_pdf_text(temp_pdf_path, temp_txt_path)

    assert result is True
    assert os.path.exists(temp_txt_path)

    # Verify the extracted text was saved
    with open(temp_txt_path, "r") as f:
        content = f.read()
        assert "OCR text line" in content
        assert "--- Page 1 ---" in content

    # Verify both methods were called
    mock_pypdfium2.assert_called_once_with(temp_pdf_path)
    mock_ocr.assert_called_once_with(temp_pdf_path)


@patch("pypdfium2.PdfDocument")
def test_extract_pdf_text_conversion_error(mock_pdf_doc, temp_pdf_path, temp_txt_path):
    """Test PDF text extraction with conversion error."""
    # Mock PDF document to raise an exception
    mock_pdf_doc.side_effect = Exception("PDF loading failed")

    result = extract_pdf_text(temp_pdf_path, temp_txt_path)

    assert result is False


def test_extract_pdf_text_nonexistent_file():
    """Test PDF text extraction with non-existent file."""
    nonexistent_path = "/path/to/nonexistent.pdf"

    result = extract_pdf_text(nonexistent_path)

    assert result is False


@pytest.mark.parametrize(
    "mock_setup,expected",
    [
        ({"text": "Sample text from PDF", "exception": False}, True),
        ({"text": "", "exception": False}, None),
        ({"text": None, "exception": True}, None),
    ],
)
@patch("pypdfium2.PdfDocument")
def test_extract_text_with_pypdfium2(
    mock_pdf_doc, temp_pdf_path, mock_setup, expected
):
    """Test text extraction with pypdfium2 in various scenarios."""
    if mock_setup["exception"]:
        mock_pdf_doc.side_effect = Exception("PDF reading failed")
    else:
        # Mock pypdfium2 PDF document and pages
        mock_pdf = MagicMock()
        mock_page = MagicMock()
        mock_textpage = MagicMock()
        mock_textpage.get_text_range.return_value = mock_setup["text"]
        mock_page.get_textpage.return_value = mock_textpage
        mock_pdf.__getitem__.return_value = mock_page
        mock_pdf.__len__.return_value = 1
        mock_pdf_doc.return_value = mock_pdf

    result = _extract_text_with_pypdfium2(temp_pdf_path)

    if expected is True:
        assert result is not None
        assert mock_setup["text"] in result
        assert "--- Page 1 ---" in result
    else:
        assert result is None


@pytest.mark.parametrize(
    "mock_setup,expected",
    [
        ({"text": "Extracted text from OCR", "conversion_error": False}, True),
        ({"text": "", "conversion_error": False}, None),
        ({"text": None, "conversion_error": True}, None),
    ],
)
@patch("pytesseract.image_to_string")
@patch("pypdfium2.PdfDocument")
def test_extract_text_with_ocr(
    mock_pdf_doc, mock_tesseract, temp_pdf_path, mock_setup, expected
):
    """Test text extraction with OCR using pypdfium2 rendering."""
    if mock_setup["conversion_error"]:
        mock_pdf_doc.side_effect = Exception("PDF loading failed")
    else:
        # Mock pypdfium2 PDF document, page, textpage, and bitmap
        mock_pdf = MagicMock()
        mock_pdf.__len__.return_value = 1
        
        mock_page = MagicMock()
        mock_textpage = MagicMock()
        mock_textpage.get_text_range.return_value = ""  # No extractable text, needs OCR
        mock_page.get_textpage.return_value = mock_textpage
        
        mock_bitmap = MagicMock()
        mock_pil_image = MagicMock()
        mock_bitmap.to_pil.return_value = mock_pil_image
        mock_page.render.return_value = mock_bitmap
        
        mock_pdf.__getitem__.return_value = mock_page
        mock_pdf_doc.return_value = mock_pdf
        
        mock_tesseract.return_value = mock_setup["text"]

    result = _extract_text_with_ocr(temp_pdf_path)

    if expected is True:
        assert result is not None
        assert mock_setup["text"] in result
        assert "--- Page 1 ---" in result
    else:
        assert result is None


def test_has_sufficient_text_enough_content():
    """Test text quality check with sufficient content."""
    text = "--- Page 1 ---\n" + "\n".join([f"Line {i}" for i in range(10)])
    assert _has_sufficient_text(text) is True


def test_has_sufficient_text_sparse_content():
    """Test text quality check with sparse content (< 5 lines per page)."""
    text = "--- Page 1 ---\nLine 1\nLine 2\n"
    assert _has_sufficient_text(text) is False


def test_has_sufficient_text_multiple_pages():
    """Test text quality check with multiple pages."""
    # Both pages have >= 5 lines (sufficient)
    text = (
        "--- Page 1 ---\n"
        + "\n".join([f"Line {i}" for i in range(10)])
        + "\n--- Page 2 ---\n"
        + "\n".join([f"Line {i}" for i in range(10)])
    )
    assert _has_sufficient_text(text) is True

    # Both pages have < 5 lines (insufficient)
    text = "--- Page 1 ---\nLine 1\nLine 2\nLine 3\nLine 4\n--- Page 2 ---\nLine 5\nLine 6\nLine 7\nLine 8\n"
    assert _has_sufficient_text(text) is False


def test_has_sufficient_text_mixed_pages():
    """Test text quality check when only some pages have sparse content."""
    # Page 1 has 10 lines (good), but Page 2 has only 2 lines (sparse)
    # With 2 pages, max_sparse_allowed = max(1, 2//5) = 1, so 1 sparse page is OK
    text = (
        "--- Page 1 ---\n"
        + "\n".join([f"Line {i}" for i in range(10)])
        + "\n--- Page 2 ---\nLine 1\nLine 2\n"
    )
    assert _has_sufficient_text(text) is True

    # Page 1 has 2 lines (sparse), Page 2 has 10 lines (good)
    # With 2 pages, 1 sparse page is allowed (e.g., signature page at end)
    text = "--- Page 1 ---\nLine 1\nLine 2\n--- Page 2 ---\n" + "\n".join(
        [f"Line {i}" for i in range(10)]
    )
    assert _has_sufficient_text(text) is True


def test_has_sufficient_text_too_many_sparse_pages():
    """Test text quality check when too many pages have sparse content."""
    # 3 pages with 2 sparse - max_sparse_allowed = max(1, 3//5) = 1, so 2 sparse fails
    text = (
        "--- Page 1 ---\n"
        + "\n".join([f"Line {i}" for i in range(10)])
        + "\n--- Page 2 ---\nLine 1\nLine 2\n"
        + "--- Page 3 ---\nLine 1\nLine 2\n"
    )
    assert _has_sufficient_text(text) is False

    # 10 pages with 3 sparse - max_sparse_allowed = max(1, 10//5) = 2, so 3 sparse fails
    text = "--- Page 1 ---\n" + "\n".join([f"Line {i}" for i in range(10)])
    for i in range(2, 8):
        text += f"\n--- Page {i} ---\n" + "\n".join([f"Line {j}" for j in range(10)])
    # Add 3 sparse pages at the end
    for i in range(8, 11):
        text += f"\n--- Page {i} ---\nLine 1\nLine 2\n"
    assert _has_sufficient_text(text) is False


def test_has_sufficient_text_allowed_sparse_pages():
    """Test text quality check with sparse pages within allowed limit."""
    # 10 pages with 2 sparse - max_sparse_allowed = max(1, 10//5) = 2, so 2 sparse is OK
    text = "--- Page 1 ---\n" + "\n".join([f"Line {i}" for i in range(10)])
    for i in range(2, 9):
        text += f"\n--- Page {i} ---\n" + "\n".join([f"Line {j}" for j in range(10)])
    # Add 2 sparse pages at the end
    for i in range(9, 11):
        text += f"\n--- Page {i} ---\nLine 1\nLine 2\n"
    assert _has_sufficient_text(text) is True


def test_has_sufficient_text_empty():
    """Test text quality check with empty text."""
    assert _has_sufficient_text("") is False
    assert _has_sufficient_text(None) is False


@patch("pydocai.extractor._extract_text_with_ocr")
@patch("pydocai.extractor._extract_text_with_pypdfium2")
def test_extract_pdf_text_sparse_fallback_to_ocr(
    mock_pypdfium2, mock_ocr, temp_pdf_path, temp_txt_path
):
    """Test PDF text extraction falls back to OCR when text is sparse."""
    # Mock pypdfium2 to return sparse text (< 5 lines per page)
    sparse_text = "--- Page 1 ---\nLine 1\nLine 2\n"
    mock_pypdfium2.return_value = sparse_text

    # Mock OCR to return rich text
    rich_ocr_text = "--- Page 1 ---\n" + "\n".join([f"OCR Line {i}" for i in range(20)])
    mock_ocr.return_value = rich_ocr_text

    result = extract_pdf_text(temp_pdf_path, temp_txt_path)

    assert result is True
    assert os.path.exists(temp_txt_path)

    # Verify the OCR text was saved (not the sparse text)
    with open(temp_txt_path, "r") as f:
        content = f.read()
        assert "OCR Line" in content

    # Verify both methods were called
    mock_pypdfium2.assert_called_once_with(temp_pdf_path)
    mock_ocr.assert_called_once_with(temp_pdf_path)
