"""LLM 信息提取模块单元测试

覆盖范围：
- _parse_json_response：纯 JSON、Markdown 包裹、带前缀、解析失败
- _image_to_base64：图片转 base64
- extract_document：成功提取、LLM 失败
"""

from __future__ import annotations

import base64
import json
from pathlib import Path

import pytest

from src.doc.llm_extractor import (
    _image_to_base64,
    _parse_json_response,
    extract_document,
)

# 字段键名（与源码中的字符串字面量保持一致）
K_INVOICE_NUMBER = "invoice_number"
K_TOTAL_AMOUNT = "total_amount"
K_CARD_DATE = "card_date"
K_CARD_NO = "card_no"
K_CARD_AMOUNT = "card_amount"

# ------------------------------------------------------------------
# _parse_json_response
# ------------------------------------------------------------------


class TestParseJsonResponse:
    """JSON 响应解析"""

    def test_pure_json(self):
        raw = json.dumps({K_INVOICE_NUMBER: "123456", K_TOTAL_AMOUNT: "100.00"})
        result = _parse_json_response(raw)
        assert result[K_INVOICE_NUMBER] == "123456"
        assert result[K_TOTAL_AMOUNT] == "100.00"

    def test_markdown_json_block(self):
        raw = f'```json\n{{"{K_INVOICE_NUMBER}": "789"}}\n```'
        result = _parse_json_response(raw)
        assert result[K_INVOICE_NUMBER] == "789"

    def test_markdown_block_without_lang(self):
        raw = '```\n{"key": "value"}\n```'
        result = _parse_json_response(raw)
        assert result["key"] == "value"

    def test_json_prefix(self):
        raw = f'json\n{{"{K_INVOICE_NUMBER}": "001"}}'
        result = _parse_json_response(raw)
        assert result[K_INVOICE_NUMBER] == "001"

    def test_json_prefix_with_whitespace(self):
        raw = '  json  \n{"a": 1}'
        result = _parse_json_response(raw)
        assert result["a"] == 1

    def test_nested_json(self):
        raw = json.dumps({"outer": {"inner": [1, 2, 3]}})
        result = _parse_json_response(raw)
        assert result["outer"]["inner"] == [1, 2, 3]

    def test_whitespace_around_json(self):
        raw = '  \n  {"x": 42}  \n  '
        result = _parse_json_response(raw)
        assert result["x"] == 42

    def test_invalid_json_raises(self):
        with pytest.raises(ValueError):
            _parse_json_response("not json at all")

    def test_empty_string_raises(self):
        with pytest.raises(ValueError):
            _parse_json_response("")


# ------------------------------------------------------------------
# _image_to_base64
# ------------------------------------------------------------------


class TestImageToBase64:
    """图片转 base64"""

    def test_png_to_base64(self, tmp_path: Path):
        content = b"\x89PNG\r\n\x1a\nfake_png_data"
        img_path = tmp_path / "test.png"
        img_path.write_bytes(content)

        result = _image_to_base64(img_path)
        assert isinstance(result, str)
        assert base64.b64decode(result) == content

    def test_jpg_to_base64(self, tmp_path: Path):
        content = b"\xff\xd8\xff\xe0fake_jpg_data"
        img_path = tmp_path / "test.jpg"
        img_path.write_bytes(content)

        result = _image_to_base64(img_path)
        assert base64.b64decode(result) == content

    def test_file_not_found_raises(self, tmp_path: Path):
        img_path = tmp_path / "nonexistent.png"
        with pytest.raises(FileNotFoundError):
            _image_to_base64(img_path)

    def test_returns_utf8_string(self, tmp_path: Path):
        content = b"test_image_content"
        img_path = tmp_path / "test.png"
        img_path.write_bytes(content)

        result = _image_to_base64(img_path)
        assert type(result) is str


# ------------------------------------------------------------------
# extract_document (mock LLM)
# ------------------------------------------------------------------


class TestExtractDocument:
    """图片提取：通过 mock _llm_query_multimodal 避免真实 LLM 调用"""

    def _mock_multimodal(self, monkeypatch, response_text: str):
        def fake_query(system_prompt, text, image_b64, max_tokens=4096):
            return response_text

        monkeypatch.setattr("src.doc.llm_extractor._llm_query_multimodal", fake_query)

    def test_success(self, tmp_path: Path, monkeypatch):
        img_path = tmp_path / "card.png"
        img_path.write_bytes(b"fake_image")

        mock_result = {
            K_CARD_DATE: "2026-01-10",
            K_CARD_NO: "6228480000000000",
            K_CARD_AMOUNT: "500.00",
        }
        self._mock_multimodal(monkeypatch, json.dumps(mock_result))

        result = extract_document(img_path)
        assert result[K_CARD_DATE] == "2026-01-10"
        assert result[K_CARD_NO] == "6228480000000000"
        assert result[K_CARD_AMOUNT] == "500.00"

    def test_llm_failure_propagates(self, tmp_path: Path, monkeypatch):
        img_path = tmp_path / "card.png"
        img_path.write_bytes(b"fake_image")

        def fake_query(system_prompt, text, image_b64, max_tokens=4096):
            raise RuntimeError("模型不可用")

        monkeypatch.setattr("src.doc.llm_extractor._llm_query_multimodal", fake_query)
        with pytest.raises(RuntimeError, match="模型不可用"):
            extract_document(img_path)

    def test_file_not_found_propagates(self, tmp_path: Path, monkeypatch):
        img_path = tmp_path / "missing.png"
        self._mock_multimodal(monkeypatch, "{}")
        with pytest.raises(FileNotFoundError):
            extract_document(img_path)

    def test_markdown_wrapped_json(self, tmp_path: Path, monkeypatch):
        img_path = tmp_path / "card.png"
        img_path.write_bytes(b"fake_image")

        mock_result = {
            K_CARD_DATE: "2026-02-01",
            K_CARD_NO: "6228481111111111",
            K_CARD_AMOUNT: "300.00",
        }
        wrapped = f"```json\n{json.dumps(mock_result)}\n```"
        self._mock_multimodal(monkeypatch, wrapped)

        result = extract_document(img_path)
        assert result[K_CARD_DATE] == "2026-02-01"

    def test_image_encoded_as_base64(self, tmp_path: Path, monkeypatch):
        img_path = tmp_path / "card.png"
        expected_content = b"test_image_data"
        img_path.write_bytes(expected_content)

        received_b64 = None

        def capture_b64(system_prompt, text, image_b64s, max_tokens=4096):
            nonlocal received_b64
            received_b64 = image_b64s
            return json.dumps({K_CARD_DATE: "2026-01-01", K_CARD_NO: "0000", K_CARD_AMOUNT: "100"})

        monkeypatch.setattr("src.doc.llm_extractor._llm_query_multimodal", capture_b64)
        extract_document(img_path)

        assert received_b64 is not None
        assert isinstance(received_b64, list)
        assert len(received_b64) >= 1
        assert base64.b64decode(received_b64[0]) == expected_content
