import time as time_module

from pydantic import BaseModel, Field, ValidationError

from ..exceptions import OcrError, OcrFailureCode
from ..models import OCRResult
from ..parse_ocr.wtp_left_page_parser import WTPLeftPageOCR, WTPLeftPageParser
from .client import DEFAULT_GEMINI_MODEL, generate_structured_content
from .convert import convert_model


class GeminiSandFilterEntry(BaseModel):
    time: str | None = None
    line_number: str | None = None
    appearance_and_odour: str | None = None
    p_alkalinity: str | None = None
    m_alkalinity: str | None = None
    free_chlorine: str | None = None
    pH: str | None = None
    turbidity: str | None = None


class GeminiUfEntry(BaseModel):
    time: str | None = None
    turbidity: str | None = None


class GeminiAcfEntry(BaseModel):
    time: str | None = None
    line_number: str | None = None
    appearance_and_odour: str | None = None
    m_alkalinity: str | None = None
    total_hardness: str | None = None
    total_chlorine: str | None = None
    pH: str | None = None
    turbidity: str | None = None


class GeminiLeadEntry(BaseModel):
    time: str | None = None
    appearance_and_odour: str | None = None
    pH: str | None = None
    tds: str | None = None
    turbidity: str | None = None


class GeminiLagEntry(BaseModel):
    time: str | None = None
    line_number: str | None = None
    appearance_and_odour: str | None = None
    pH: str | None = None
    tds: str | None = None
    turbidity: str | None = None


class GeminiRoEntry(BaseModel):
    time: str | None = None
    appearance_odour_taste: str | None = None
    feed_pH: str | None = None
    feed_tds: str | None = None
    permeate_pH: str | None = None
    permeate_tds: str | None = None
    turbidity: str | None = None


class GeminiWTPLeftPageOCR(BaseModel):
    date: str | None = None
    sand_filter_entries: list[GeminiSandFilterEntry] = Field(default_factory=list)
    uf_entries: list[GeminiUfEntry] = Field(default_factory=list)
    lead_entries: list[GeminiLeadEntry] = Field(default_factory=list)
    lag_entries: list[GeminiLagEntry] = Field(default_factory=list)
    ro_entries: list[GeminiRoEntry] = Field(default_factory=list)
    acf_entries: list[GeminiAcfEntry] = Field(default_factory=list)


WTP_LEFT_PAGE_PROMPT = """
You are transcribing a single photographed page from a handwritten Water Treatment Plant (WTP) daily
monitoring logbook ("Enrich Agro Food Products Pvt. Ltd. Rohtak", form QSE-8.1-FRM-WTP-1.1.1.10).
The page contains one date field and 6 handwritten data tables, each a series of rows keyed by time.
Your ONLY job is literal, verbatim transcription of what is written in each cell — do not calculate,
round, unit-convert, or "correct" any value.

Page layout overview: the six tables are arranged in a 3x2 grid, top to bottom, left to right:
- Top row: Sand Filter (left) | UF (right, next to Sand Filter).
- Middle row: ACF (left, below Sand Filter/UF) | Lead (right, next to ACF).
- Bottom row: Lag (left, below ACF) | RO (right, next to Lag).
Each table is its own independent list of rows, column-bound only by its own vertical rule lines —
never borrow a time or value from whichever table happens to sit alongside it on the same
horizontal line, and never stop transcribing one table just because a neighboring table above,
below, or beside it has ended.

Transcribe the following, using EXACTLY the JSON field names given below:

1. date: the handwritten date in the "Date" box at the top-left of the page, exactly as written
   (e.g. "17/4/26"). null if blank or illegible.

2. sand_filter_entries: rows of the "Sand Filter" table (top-left). Columns, left to right: time,
   line_number ("Line No."), appearance_and_odour ("Appearance and Odor"), p_alkalinity
   ("p-Alkalinity"), m_alkalinity ("M Alkalinity"), free_chlorine ("Free Chlorine"), pH, turbidity.

3. uf_entries: rows of the "UF" table (top-right, next to Sand Filter). Columns: time, turbidity.

4. acf_entries: rows of the "ACF" table (below Sand Filter/UF). Columns, left to right: time,
   line_number ("Line No."), appearance_and_odour ("Appearance and Odor"), m_alkalinity
   ("M Alkalinity"), total_hardness ("Total Hardness"), total_chlorine ("Total Chlorine"), pH,
   turbidity.

5. lead_entries: rows of the "Lead" table (next to ACF). Columns, left to right: time,
   appearance_and_odour ("Appearance and Odor"), pH, turbidity, tds ("TDS").

6. lag_entries: rows of the "Lag" table (below ACF). Columns, left to right: time,
   line_number ("Line No."), appearance_and_odour ("Appearance and Odor"), pH, turbidity,
   tds ("TDS").

7. ro_entries: rows of the "RO" table (next to Lag). Columns, left to right: time,
   appearance_odour_taste ("Appearance, Odor & Taste"), feed_pH ("Feed pH"), permeate_pH
   ("Permeate pH"), feed_tds ("Feed TDS"), permeate_tds ("Permeate TDS"), turbidity.

Ignore the "SIGNATURE OF CHEMIST" / "SIGNATURE OF MANAGER" section and any printed boilerplate
(form code, revision info) at the bottom of the page — those are not part of the JSON output.

General rules for every table:
- Transcribe each handwritten cell's text EXACTLY as written (same digits, abbreviations,
  punctuation, e.g. "Nm|Noo" or "03 P2"). Do not compute, round, normalize units, or infer values
  that are not written.
- Use null for any cell that is blank, illegible, or crossed out.
- Skip (do not emit) a row that is entirely empty.
- Only emit as many rows per table as are actually filled in; do not invent or pad rows.
- If a table has no filled rows at all, return an empty list [] for it.
- Respond with ONLY the JSON object matching the given response schema — no extra commentary.
"""


def get_wtp_left_page_ocr_gemini(
    file_path: str,
    mime_type: str = "image/jpeg",
    model: str = DEFAULT_GEMINI_MODEL,
    raw_gemini_output_path: str | None = None,
) -> OCRResult[WTPLeftPageOCR]:
    start = time_module.perf_counter()

    raw_result = generate_structured_content(
        file_path=file_path,
        mime_type=mime_type,
        prompt=WTP_LEFT_PAGE_PROMPT,
        response_schema=GeminiWTPLeftPageOCR,
        model=model,
    )

    if raw_gemini_output_path is not None:
        with open(raw_gemini_output_path, "w") as f:
            f.write(raw_result.model_dump_json(indent=2))

    parser = WTPLeftPageParser()
    try:
        result = convert_model(parser, raw_result, WTPLeftPageOCR)
    except ValidationError as exc:
        raise OcrError(
            OcrFailureCode.EXTRACTION_VALIDATION_FAILED,
            f"Extracted data failed validation: {exc}",
        ) from exc
    # log_rows is a derived cross-table field (zips same-index rows across the 6 tables and
    # flags any whose times aren't sequential/within 15 min of each other) — not something
    # Gemini can produce directly, so it's rebuilt with the same logic the Document AI path uses.
    result.log_rows = parser._build_log_rows(result)

    return OCRResult(
        result=result,
        time_elapsed=time_module.perf_counter() - start,
    )
