from __future__ import annotations

import time
from dataclasses import dataclass

from . import __version__
from .ollama import OllamaClient, OllamaError
from .prompts import transcription_prompt


@dataclass(frozen=True)
class InputDocument:
    document_id: str
    name: str
    document_type: str
    content: bytes


class Engine:
    def __init__(self, client: OllamaClient) -> None:
        self.client = client

    def extract(self, documents: list[InputDocument]) -> dict:
        started = time.perf_counter()
        output_documents: list[dict] = []
        transcriptions: list[dict] = []
        warnings: list[str] = []

        for document in documents:
            item_started = time.perf_counter()
            text = ""
            error_code = ""
            try:
                result = self.client.transcribe(
                    document.content,
                    transcription_prompt(document.document_type, self.client.model),
                    document.document_type,
                )
                text = result.text
                eval_count = result.eval_count
                model_duration_ms = round(result.total_duration_ns / 1_000_000)
            except OllamaError:
                eval_count = 0
                model_duration_ms = 0
                error_code = "ollama_error"
                warnings.append(f"{document.name}: a VLM nao conseguiu transcrever esta pagina")

            duration_ms = round((time.perf_counter() - item_started) * 1000)
            output_documents.append(
                {
                    "id": document.document_id,
                    "name": document.name,
                    "type": document.document_type or "desconhecido",
                    "ocr_ok": bool(text),
                    "chars": len(text),
                    "duration_ms": duration_ms,
                    "model_duration_ms": model_duration_ms,
                    "eval_count": eval_count,
                    "error_code": error_code,
                }
            )
            transcriptions.append(
                {
                    "document_id": document.document_id,
                    "name": document.name,
                    "type_hint": document.document_type,
                    "text": text,
                }
            )

        return {
            "success": not any(item["error_code"] for item in output_documents),
            "engine_version": f"pataniscAI-{__version__}",
            "model": self.client.model,
            "duration_ms": round((time.perf_counter() - started) * 1000),
            "documents": output_documents,
            "transcriptions": transcriptions,
            "results": {},
            "avisos": warnings,
            "incertos": [],
        }