from typing import Literal

from pydantic import BaseModel

from .client import generate_structured_content

DEFAULT_QUALITY_CHECK_MODEL = "gemini-2.5-flash-lite"
# 0 disables thinking entirely (valid for flash/flash-lite, unlike pro which has a thinking floor)
# -- this is a coarse pass/fail check, not worth spending thinking tokens on.
DEFAULT_QUALITY_CHECK_THINKING_BUDGET = 0

DocumentQualityVerdict = Literal["ACCEPTED", "REJECTED"]


class _DocumentQualityResponse(BaseModel):
    verdict: DocumentQualityVerdict
    reason: str | None = None


DOCUMENT_QUALITY_PROMPT = """
You are a fast quality gate for photographs of handwritten Water Treatment Plant (WTP) logbook
pages, before they are sent on for detailed data extraction. Look at the attached image and decide
whether it is usable for that purpose.

Respond REJECTED if any of the following are true:
- The image is blurry, out of focus, or motion-smeared enough that handwritten numbers or letters
  in the data tables cannot be confidently read.
- The photo is cropped, cut off, or angled such that part of a data table (rows, columns, or
  header) is missing or outside the frame.
- Lighting, glare, shadow, or obstruction (e.g. a hand, finger, or object) covers enough of the
  page that a meaningful portion of the handwritten data is unreadable.

Respond ACCEPTED if the page's tables and handwritten entries are legible and fully in frame, even
if there is some minor glare, a normal amount of handwriting messiness, or an incidental object
(e.g. a fingertip) at the very edge that does not obscure any data.

Give a one-sentence reason for your verdict. Respond with ONLY the JSON object matching the given
response schema — no extra commentary.
"""


def check_document_quality(
    file_path: str,
    mime_type: str = "image/jpeg",
    model: str = DEFAULT_QUALITY_CHECK_MODEL,
    thinking_budget: int = DEFAULT_QUALITY_CHECK_THINKING_BUDGET,
) -> DocumentQualityVerdict:
    result = generate_structured_content(
        file_path=file_path,
        mime_type=mime_type,
        prompt=DOCUMENT_QUALITY_PROMPT,
        response_schema=_DocumentQualityResponse,
        model=model,
        thinking_budget=thinking_budget,
    )

    print(f"[gemini] document quality verdict={result.verdict} reason={result.reason!r}")
    return result.verdict
