# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

import base64
import json
from io import BytesIO

import httpx
from PIL import Image
import pytest

from nemotron_parse.inference import (
    DocumentParser,
    GenerationError,
    GenerationSettings,
    merge_crops,
)
from nemotron_parse import PostprocessOptions
from test_postprocessing import block, tabular

OLD = tabular(r"\thead{Name} & \thead{Value}", "A & old")
NEW = tabular(r"\thead{Name} & \thead{Value}", "A & new")


def test_double_refines_table_and_keeps_text_order_and_geometry():
    calls = []

    def backend(image, *, stage, **kwargs):
        calls.append((stage, image.size))
        return {
            "raw_out": (
                (
                    block("Before", box=(0, 0, 1, 0.05))
                    + block(OLD, "Table")
                    + block("After", box=(0, 0.95, 1, 1))
                )
                if stage == "fullpage"
                else block(NEW, "Table", (0, 0, 1, 1))
            ),
            "finish_reason": "stop",
        }

    with DocumentParser(
        backend=backend, crop_margin_px=0, crop_margin_ratio=0, merge_policy="guarded"
    ) as parser:
        doc = parser.parse_double(Image.new("RGB", (100, 200)))
    assert [c[0] for c in calls] == ["fullpage", "crop"]
    assert calls[1][1] == (80, 160)
    assert doc["inference"]["requests"] == 2
    assert [b["class"] for b in doc["blocks"]] == ["Text", "Table", "Text"]
    assert doc["blocks"][1]["bbox"] == [0.1, 0.1, 0.9, 0.9]
    assert "new" in doc["markdown"] and "old" not in doc["markdown"]
    assert "old" in doc["stages"]["fullpage"]["raw_out"]


def test_triple_detects_missing_table_and_maps_crop_coordinates():
    def backend(image, *, stage, **kwargs):
        if stage == "fullpage":
            return block("Before", box=(0, 0, 1, 0.05))
        if stage == "detection":
            return block("", "Table", (0.2, 0.3, 0.8, 0.7))
        return block(NEW, "Table", (0, 0, 1, 1))

    parser = DocumentParser(backend=backend, crop_margin_px=0, crop_margin_ratio=0)
    doc = parser.parse_triple(Image.new("RGB", (100, 100)))
    assert doc["inference"]["requests"] == 3
    assert doc["blocks"][1]["bbox"] == [0.2, 0.3, 0.8, 0.7]
    assert "new" in doc["markdown"]


def test_union_adds_parser_regions_but_detector_only_does_not():
    def backend(image, *, stage, **kwargs):
        if stage == "fullpage":
            return block(OLD, "Table", (0.1, 0.1, 0.4, 0.4))
        if stage == "detection":
            return block("", "Chart", (0.6, 0.6, 0.9, 0.9))
        return block("chart data", "Chart", (0, 0, 1, 1))

    parser = DocumentParser(backend=backend, crop_margin_px=0, crop_margin_ratio=0)
    image = Image.new("RGB", (100, 100))
    union = parser.parse_triple(image)
    detector = parser.parse_triple(image, triple_regions="detector")
    assert len(union["stages"]["crops"]) == 2
    assert len(detector["stages"]["crops"]) == 1


def test_region_limit_applies_to_union():
    def backend(image, *, stage, **kwargs):
        return (
            block("", "Table", (0.1, 0.1, 0.4, 0.4))
            if stage == "detection"
            else block(OLD, "Table", (0.6, 0.6, 0.9, 0.9))
        )

    parser = DocumentParser(backend=backend, max_regions=1)
    assert (
        len(parser.parse_triple(Image.new("RGB", (100, 100)))["inference"]["regions"])
        == 1
    )


def test_no_regions_does_not_make_crop_request():
    calls = []

    def backend(image, *, stage, **kwargs):
        calls.append(stage)
        return block("hello") if stage == "fullpage" else "</s>"

    parser = DocumentParser(backend=backend)
    doc = parser.parse_triple(Image.new("RGB", (20, 20)))
    assert calls == ["fullpage", "detection"]
    assert doc["markdown"] == "hello"


@pytest.mark.parametrize("stage", ["fullpage", "detection", "crop"])
@pytest.mark.parametrize(
    "response",
    [
        {"raw_out": "partial", "finish_reason": "length"},
        {"raw_out": ""},
        {"raw_out": "bad", "error": "failed"},
    ],
)
def test_failed_generations_are_not_reported_as_success(stage, response):
    def backend(image, *, stage: str, **kwargs):
        if stage == target:
            return response
        return block(OLD, "Table")

    target = stage
    parser = DocumentParser(backend=backend)
    with pytest.raises(GenerationError) as error:
        parser.parse_triple(Image.new("RGB", (100, 100)))
    assert error.value.stage == stage
    assert error.value.response["raw_out"] == response["raw_out"]


def test_guard_rejects_incompatible_column_count():
    wrong = tabular(
        r"\thead{Name} & \thead{Value} & \thead{Extra}", "A & new & extra", cols="ccc"
    )
    row = {
        "raw_out": block(wrong, "Table", (0, 0, 1, 1)),
        "detector_label": "table",
        "detector_score": 1.0,
        "detector_bbox_norm": (0.1, 0.1, 0.9, 0.9),
        "crop_bbox_pixels": (10, 10, 90, 90),
        "page_size": (100, 100),
    }
    merged = merge_crops(block(OLD, "Table"), [row], policy="guarded")
    assert "old" in merged and "extra" not in merged


def test_batch_preserves_order(tmp_path):
    paths = []
    for width in [10, 20, 30]:
        path = tmp_path / f"{width}.png"
        Image.new("RGB", (width, 10)).save(path)
        paths.append(path)
    parser = DocumentParser(backend=lambda image, **kwargs: block(str(image.width)))
    docs = parser.parse_many(paths, concurrency=3)
    assert [d["markdown"] for d in docs] == ["10", "20", "30"]


def test_http_endpoint_protocol_auth_and_stage_settings():
    calls = []

    def handle(request):
        body = json.loads(request.content)
        calls.append(body)
        assert request.headers["Authorization"] == "Bearer test-key"
        assert str(request.url) == "http://example/v1/chat/completions"
        assert body["skip_special_tokens"] is False
        uri = body["messages"][0]["content"][1]["image_url"]["url"]
        image = Image.open(BytesIO(base64.b64decode(uri.split(",")[1])))
        assert image.mode == "RGB"
        prompt = body["messages"][0]["content"][0]["text"]
        raw = block("", "Table") if "output_no_text" in prompt else block(OLD, "Table")
        return httpx.Response(
            200,
            json={
                "choices": [{"message": {"content": raw}, "finish_reason": "stop"}],
                "usage": {"completion_tokens": 10},
            },
        )

    client = httpx.Client(transport=httpx.MockTransport(handle))
    with DocumentParser(
        "http://example/v1/", api_key="test-key", client=client
    ) as parser:
        doc = parser.parse_triple(Image.new("RGB", (100, 100)))
    assert [x["max_tokens"] for x in calls] == [8744, 4096, 8744]
    assert [x["temperature"] for x in calls] == [0.2, 0.2, 0.1]
    assert doc["stages"]["fullpage"]["usage"]["completion_tokens"] == 10
    assert not client.is_closed
    client.close()


@pytest.mark.parametrize("mode", ["bad", "all"])
def test_invalid_mode_does_not_issue_request(mode):
    parser = DocumentParser(backend=lambda *a, **k: pytest.fail("unexpected inference"))
    with pytest.raises(ValueError):
        parser.parse(Image.new("RGB", (10, 10)), mode=mode)


def test_generation_options_validate_limits():
    with pytest.raises(ValueError):
        GenerationSettings(max_tokens=0)
    with pytest.raises(ValueError):
        GenerationSettings(temperature=float("nan"))
    with pytest.raises(ValueError):
        DocumentParser(backend=lambda **k: "", max_regions=0)


def test_transformers_backend_cpu_settings_rng_and_truncation():
    import torch
    from types import SimpleNamespace
    from nemotron_parse.inference import TransformersBackend

    class Inputs(dict):
        def to(self, device):
            return self

    class Processor:
        def __call__(self, **kwargs):
            assert kwargs["add_special_tokens"] is False
            return Inputs(input_ids=torch.tensor([[4, 5]]))

        def batch_decode(self, sequences, **kwargs):
            assert kwargs["skip_special_tokens"] is False
            return [block("hello")]

    class Model:
        device = "cpu"
        generation_config = SimpleNamespace(max_new_tokens=999)

        def generate(self, **kwargs):
            config = kwargs["generation_config"]
            assert config.do_sample is False and config.max_new_tokens == 3
            torch.rand(1)
            return SimpleNamespace(sequences=torch.tensor([[4, 5, 7, 2]]))

    model = Model()
    backend = TransformersBackend(model, Processor())
    before = torch.random.get_rng_state().clone()
    response = backend(
        Image.new("RGB", (10, 10)),
        prompt="test",
        stage="fullpage",
        settings=GenerationSettings(temperature=0, max_tokens=3),
    )
    assert response["usage"]["completion_tokens"] == 2
    assert response["finish_reason"] == "stop"
    assert torch.equal(before, torch.random.get_rng_state())
    assert model.generation_config.max_new_tokens == 999

    class Capped(Model):
        def generate(self, **kwargs):
            return SimpleNamespace(sequences=torch.tensor([[4, 5, 7, 8, 2]]))

    backend = TransformersBackend(Capped(), Processor())
    parser = DocumentParser(
        backend=backend, fullpage=GenerationSettings(temperature=0, max_tokens=3)
    )
    with pytest.raises(GenerationError):
        parser.parse(Image.new("RGB", (10, 10)))


def test_chart_crop_refinement_is_guarded():
    original = "<chart>\ntitle: Sales\n| Item | Value |\n| --- | --- |\n| A | 10 |\n| B | 20 |\n</chart>"
    updated = original.replace("10", "11")
    row = {
        "raw_out": block(updated, "Chart", (0, 0, 1, 1)),
        "detector_label": "chart",
        "detector_score": 1.0,
        "detector_bbox_norm": (0.1, 0.1, 0.9, 0.9),
        "crop_bbox_pixels": (10, 10, 90, 90),
        "page_size": (100, 100),
    }
    merged = merge_crops(block(original, "Chart"), [row], policy="guarded")
    assert merged.count("<class_Chart>") == 1 and "| A | 11 |" in merged
    row["raw_out"] = block(updated.replace("</chart>", ""), "Chart", (0, 0, 1, 1))
    merged = merge_crops(block(original, "Chart"), [row], policy="guarded")
    assert "| A | 10 |" in merged and merged.count("<class_Chart>") == 1


def test_cli_writes_markdown_and_stage_json(tmp_path, monkeypatch, capsys):
    from nemotron_parse import inference

    image = tmp_path / "page.png"
    Image.new("RGB", (20, 20)).save(image)
    output = tmp_path / "results"
    original = inference.DocumentParser
    monkeypatch.setattr(
        inference,
        "DocumentParser",
        lambda *a, **k: original(
            backend=lambda image, **kw: block("CLI text"), options=k.get("options")
        ),
    )
    monkeypatch.setattr(
        "sys.argv",
        [
            "nemotron-parse",
            "--image",
            str(image),
            "--mode",
            "double",
            "--write-page-files",
            "--output",
            str(output),
        ],
    )
    inference.main()
    assert (output / "0000_page.md").read_text() == "CLI text"
    saved = json.loads((output / "0000_page.json").read_text())
    assert saved["inference"]["mode"] == "double"
    assert saved["stages"]["fullpage"]["raw_out"] == block("CLI text")
    with pytest.raises(SystemExit):
        inference.main()


def test_default_merge_retains_disagreeing_chart_alternatives():
    original = "<chart>\n| Category | Value |\n| --- | --- |\n| Apples | 10 |\n| Pears | 20 |\n</chart>"
    alternate = "<chart>\n| Category | Value |\n| --- | --- |\n| Apples | 15 |\n| Pears | 25 |\n</chart>"
    row = {
        "raw_out": block(alternate, "Chart", (0, 0, 1, 1)),
        "detector_label": "chart",
        "detector_score": 1.0,
        "detector_bbox_norm": (0.1, 0.1, 0.9, 0.9),
        "crop_bbox_pixels": (10, 10, 90, 90),
        "page_size": (100, 100),
    }
    result = merge_crops(block(original, "Chart"), [row])
    assert original in result and alternate in result
    assert result.count("<class_Chart>") == 2


def test_default_merge_deduplicates_matching_chart_alternatives():
    original = "<chart>\n| Category | Value |\n| --- | --- |\n| Apples | 10 |\n| Pears | 20 |\n</chart>"
    row = {
        "raw_out": block(original, "Chart", (0, 0, 1, 1)),
        "detector_label": "chart",
        "detector_score": 1.0,
        "detector_bbox_norm": (0.1, 0.1, 0.9, 0.9),
        "crop_bbox_pixels": (10, 10, 90, 90),
        "page_size": (100, 100),
    }
    assert merge_crops(block(original, "Chart"), [row]).count("<class_Chart>") == 1


def test_merge_policy_default_and_invalid_policy():
    with DocumentParser(backend=lambda image, **kw: block("Text")) as parser:
        assert parser.merge_policy == "alternatives"
        assert (
            parser.parse(Image.new("RGB", (20, 20)))["inference"]["merge_policy"]
            == "alternatives"
        )
    with pytest.raises(ValueError, match="merge_policy"):
        DocumentParser(
            backend=lambda image, **kw: block("Text"), merge_policy="unknown"
        )