"""Regresiones de selección de modelo en trabajos ARQ."""
from __future__ import annotations

import asyncio
from contextlib import contextmanager
from types import SimpleNamespace

from worker import tasks


class FakeSession:
    def __init__(self, task: object) -> None:
        self.task = task

    def get(self, _model: object, _identifier: str, **_kwargs: object) -> object:
        return self.task

    def add(self, _value: object) -> None:
        pass

    def flush(self) -> None:
        pass


class FakeRateLimiter:
    async def acquire(self, *, estimated_tokens: int) -> None:
        assert estimated_tokens >= 2000


class FakeService:
    def __init__(self, model: str) -> None:
        self.extractor = SimpleNamespace(model=model)
        self.calls = 0

    def extract_full_from_pdf(self, *_args: object, **_kwargs: object):
        self.calls += 1
        return {"numero": "F-1"}, {}


def test_process_document_uses_requested_model(monkeypatch, tmp_path) -> None:
    task = SimpleNamespace(
        request={
            "model": "modelo-elegido",
            "fields": ["numero"],
            "structure_mode": "general",
            "doc_types": [],
        },
        result={"processed": 0, "total": 1, "row_count": 0},
        updated_at=None,
    )
    session = FakeSession(task)

    @contextmanager
    def fake_session_scope():
        yield session

    base_service = FakeService("modelo-global")
    selected_service = FakeService("modelo-elegido")
    selected_models: list[str] = []

    def fake_service_for_model(_base: object, model: str):
        selected_models.append(model)
        return selected_service

    monkeypatch.setattr(tasks, "session_scope", fake_session_scope)
    monkeypatch.setattr(tasks, "get_rate_limiter", lambda: FakeRateLimiter())
    monkeypatch.setattr(tasks, "_service_for_model", fake_service_for_model)
    monkeypatch.setattr(tasks, "normalize_to_schema", lambda data, **_kwargs: {"normalized": data})
    monkeypatch.setattr(tasks, "flatten_extraction_result", lambda *_args, **_kwargs: [])
    monkeypatch.setattr(tasks, "flag_modified", lambda *_args, **_kwargs: None)
    monkeypatch.setattr(tasks, "mark_task_completed_if_done", lambda *_args, **_kwargs: None)

    queue_file = tmp_path / "factura.pdf"
    queue_file.write_bytes(b"pdf")
    ctx = {"service": base_service}

    asyncio.run(tasks.process_document(
        ctx,
        task_id="task-1",
        batch_id="batch-1",
        queue_path=str(queue_file),
        radicado="RAD-1",
        filename="factura.pdf",
        doc_type=None,
        instruction=None,
        tiff_page_range=None,
    ))

    assert selected_models == ["modelo-elegido"]
    assert selected_service.calls == 1
    assert base_service.calls == 0
    assert ctx["services_by_model"]["modelo-elegido"] is selected_service


def test_service_for_request_reuses_cached_model(monkeypatch) -> None:
    base_service = FakeService("modelo-global")
    selected_service = FakeService("modelo-elegido")
    builds = 0

    def fake_service_for_model(_base: object, _model: str):
        nonlocal builds
        builds += 1
        return selected_service

    monkeypatch.setattr(tasks, "_service_for_model", fake_service_for_model)
    ctx = {"service": base_service}
    request = {"model": "modelo-elegido"}

    first = tasks._service_for_request(ctx, request)
    second = tasks._service_for_request(ctx, request)

    assert first is selected_service
    assert second is selected_service
    assert builds == 1