from pataniscai.ollama import OllamaClient


class CapturingOllama(OllamaClient):
    def __init__(self, model: str) -> None:
        super().__init__("http://127.0.0.1:11434", model, 30)
        self.path = ""
        self.body = {}

    def _request(self, path: str, body: dict | None = None) -> dict:
        self.path = path
        self.body = body or {}
        return {"response": "texto", "eval_count": 2}


def test_paddle_keeps_generate_and_deterministic_options() -> None:
    client = CapturingOllama("AuditAid/PaddleOCR-VL-1.6-0.9B:latest")

    client.transcribe(b"image", "Transcreve.")

    assert client.path == "/api/generate"
    assert client.body["prompt"] == "Transcreve."
    assert client.body["options"] == {
        "temperature": 0,
        "num_predict": 1400,
        "repeat_penalty": 1.15,
        "repeat_last_n": 256,
    }


def test_paddle_cc_front_focus_uses_short_generation_budget() -> None:
    client = CapturingOllama("AuditAid/PaddleOCR-VL-1.6-0.9B:latest")

    client.transcribe(b"image", "Transcreve as linhas.", "cc_frente_identificacao")

    assert client.body["options"] == {
        "temperature": 0,
        "num_predict": 160,
        "repeat_penalty": 1.05,
        "repeat_last_n": 64,
    }


def test_paddle_rc_focus_uses_short_generation_budget() -> None:
    client = CapturingOllama("AuditAid/PaddleOCR-VL-1.6-0.9B:latest")

    client.transcribe(b"image", "Transcreve emissao e validade.", "rc_validade_focada")

    assert client.body["options"] == {
        "temperature": 0,
        "num_predict": 160,
        "repeat_penalty": 1.05,
        "repeat_last_n": 64,
    }


def test_glm_keeps_generate_and_canonical_options() -> None:
    client = CapturingOllama("glm-ocr:latest")

    client.transcribe(b"image", "Text Recognition:")

    assert client.path == "/api/generate"
    assert client.body["prompt"] == "Text Recognition:"
    assert client.body["options"]["temperature"] == 0
    assert client.body["options"]["top_k"] == 1
    assert client.body["options"]["num_predict"] == 2400