from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import pytest

from services.pdf_ocr_service import extract_scanned_pdf_text


class FakePixmap:
    def tobytes(self, *_args, **_kwargs):
        return b"jpeg-page"


class FakePage:
    def get_pixmap(self, **_kwargs):
        return FakePixmap()


class FakeDocument:
    page_count = 2

    def __enter__(self):
        return self

    def __exit__(self, *_args):
        return None

    def load_page(self, _index):
        return FakePage()


def test_scanned_pdf_ocr_preserves_page_order_and_reports_progress(monkeypatch):
    fake_pymupdf = SimpleNamespace(
        open=lambda _path: FakeDocument(),
        Matrix=lambda x, y: (x, y),
        csRGB="rgb",
    )
    monkeypatch.setitem(__import__("sys").modules, "pymupdf", fake_pymupdf)
    progress = []
    reservation = SimpleNamespace()
    config = SimpleNamespace(
        api_key="configured",
        ocr_model="qwen-vl-ocr",
        vision_model="vision",
        vision_max_tokens=2048,
    )

    with (
        patch("services.pdf_ocr_service.get_chat_model_config", return_value=config),
        patch("services.pdf_ocr_service.prepare_image", return_value=SimpleNamespace()),
        patch(
            "services.pdf_ocr_service.call_vision_model",
            side_effect=[
                {"content": "第一页文字", "usage": {"total_tokens": 10}},
                {"content": "第二页文字", "usage": {"total_tokens": 12}},
            ],
        ),
        patch("services.pdf_ocr_service.reserve_avatar_tokens", return_value=reservation) as reserve,
        patch("services.pdf_ocr_service.settle_reservation") as settle,
    ):
        text = extract_scanned_pdf_text(
            MagicMock(),
            SimpleNamespace(id="avatar-1"),
            "/tmp/scanned.pdf",
            on_progress=lambda done, total: progress.append((done, total)),
        )

    assert text == "[第 1 页]\n第一页文字\n\n[第 2 页]\n第二页文字"
    assert progress == [(1, 2), (2, 2)]
    assert reserve.call_count == 2
    assert settle.call_count == 2


def test_scanned_pdf_ocr_releases_tokens_after_retries_fail(monkeypatch):
    fake_document = FakeDocument()
    fake_document.page_count = 1
    fake_pymupdf = SimpleNamespace(
        open=lambda _path: fake_document,
        Matrix=lambda x, y: (x, y),
        csRGB="rgb",
    )
    monkeypatch.setitem(__import__("sys").modules, "pymupdf", fake_pymupdf)
    monkeypatch.setenv("KNOWLEDGE_PDF_OCR_ATTEMPTS", "2")
    reservation = SimpleNamespace()
    config = SimpleNamespace(
        api_key="configured",
        ocr_model="qwen-vl-ocr",
        vision_model="vision",
        vision_max_tokens=2048,
    )

    with (
        patch("services.pdf_ocr_service.get_chat_model_config", return_value=config),
        patch("services.pdf_ocr_service.prepare_image", return_value=SimpleNamespace()),
        patch("services.pdf_ocr_service.call_vision_model", side_effect=RuntimeError("timeout")) as call,
        patch("services.pdf_ocr_service.reserve_avatar_tokens", return_value=reservation),
        patch("services.pdf_ocr_service.release_reservation") as release,
        patch("services.pdf_ocr_service.time.sleep"),
    ):
        with pytest.raises(RuntimeError, match="第 1/1 页识别失败"):
            extract_scanned_pdf_text(MagicMock(), SimpleNamespace(id="avatar-1"), "/tmp/scanned.pdf")

    assert call.call_count == 2
    release.assert_called_once()