from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from ..exceptions import OcrError
from ..gemini_ocr.wtp_left_page_gemini import get_wtp_left_page_ocr_gemini
from ..gemini_ocr.wtp_right_page_gemini import get_wtp_right_page_ocr_gemini
from ..get_structured_ocr import get_wtp_coagulation_log_ocr
from .fixtures import FixtureRecord


@dataclass
class RerunResult:
    response_payload: dict[str, Any] | None
    error: str | None = None


def rerun_ocr(record: FixtureRecord) -> RerunResult:
    """Re-run OCR for one fixture record's images, producing a `response_payload`
    shaped exactly like `wra.services.ocr_service.trigger_wtp_ocr`'s output (which is
    what `expected/response_payload.json` was captured from) so the two are directly
    comparable: Gemini for RO's left/right pages, Document AI for coagulation.
    """
    try:
        if record.record_type == "ro":
            return _rerun_ro(record)
        if record.record_type == "coagulation":
            return _rerun_coagulation(record)
        return RerunResult(response_payload=None, error=f"Unsupported record_type: {record.record_type}")
    except OcrError as exc:
        return RerunResult(response_payload=None, error=f"{exc.code.value}: {exc.message}")
    except Exception as exc:
        # Not every failure mode below the library's own entrypoints is wrapped in
        # OcrError -- e.g. Document AI's client raises a raw google.auth
        # DefaultCredentialsError before it ever makes a request if Application
        # Default Credentials aren't configured. One record's environment/credential
        # problem shouldn't crash the whole batch any more than an OcrError would.
        return RerunResult(response_payload=None, error=f"{type(exc).__name__}: {exc}")


def _rerun_ro(record: FixtureRecord) -> RerunResult:
    payload: dict[str, Any] = {}

    right_image = record.images.get("right")
    if right_image is not None:
        right_result = get_wtp_right_page_ocr_gemini(str(right_image.path), right_image.content_type)
        payload["right_page"] = right_result.result.model_dump(mode="json")
        payload["right_page_time_elapsed"] = right_result.time_elapsed

    left_image = record.images.get("left")
    if left_image is not None:
        left_result = get_wtp_left_page_ocr_gemini(str(left_image.path), left_image.content_type)
        payload["left_page"] = left_result.result.model_dump(mode="json")
        payload["left_page_time_elapsed"] = left_result.time_elapsed

    if not payload:
        return RerunResult(response_payload=None, error="No left/right image in fixture record")

    return RerunResult(response_payload=payload)


def _rerun_coagulation(record: FixtureRecord) -> RerunResult:
    image = record.images.get("coagulation")
    if image is None:
        return RerunResult(response_payload=None, error="No coagulation image in fixture record")

    result = get_wtp_coagulation_log_ocr(str(image.path), image.content_type)
    return RerunResult(
        response_payload={
            "coagulation": result.result.model_dump(mode="json"),
            "coagulation_time_elapsed": result.time_elapsed,
        }
    )
