import time as time_module

from pydantic import BaseModel, Field, ValidationError

from ..exceptions import OcrError, OcrFailureCode
from ..models import OCRResult
from ..parse_ocr.wtp_right_page_parser import WTPRightPageOCR, WTPRightPageParser
from .client import DEFAULT_GEMINI_MODEL, generate_structured_content
from .convert import convert_model


class GeminiColumnFrequency(BaseModel):
    """One column's printed frequency label. Gemini's structured-output schema rejects
    open/map types (`additionalProperties`) outright, so a *_column_frequencies field
    can't be a `dict[str, str]` on this raw mirror model -- it has to be a list of these
    pairs instead. convert.py collapses the list back into a dict for the typed target
    model, where a dict is fine (it's not part of any Gemini schema)."""

    column: str
    frequency: str


class GeminiRawWaterAnalysisEntry(BaseModel):
    time: str | None = None
    raw_water_storage_tank_use: str | None = None
    appearance: str | None = None
    odour: str | None = None
    cah: str | None = None
    chloride: str | None = None
    iron: str | None = None
    m_alk: str | None = None
    pH: str | None = None
    p_alk: str | None = None
    sulphate: str | None = None
    tds: str | None = None
    th: str | None = None
    turbidity: str | None = None
    total_chlorine: str | None = None


class GeminiSoftWaterAnalysisEntry(BaseModel):
    time: str | None = None
    cah: str | None = None
    m_alk: str | None = None
    residual_chlorine: str | None = None
    taste_odour_appearance: str | None = None
    tds: str | None = None
    th: str | None = None
    turbidity: str | None = None
    pH: str | None = None
    iron: str | None = None
    chloride: str | None = None


class GeminiTreatedWaterAnalysisEntry(BaseModel):
    time: str | None = None
    aluminium: str | None = None
    appearance: str | None = None
    cah: str | None = None
    free_chlorine: str | None = None
    iron: str | None = None
    m_alkalinity: str | None = None
    odour: str | None = None
    pH: str | None = None
    p_alkalinity: str | None = None
    taste: str | None = None
    tds: str | None = None
    th: str | None = None
    total_chlorine: str | None = None
    turbidity: str | None = None
    chloride: str | None = None
    sulphate: str | None = None
    total_combined_sulphate_and_chloride: str | None = None
    remark: str | None = None


class GeminiUVLightEntry(BaseModel):
    time: str | None = None
    uv_light_number: str | None = None
    uv_light_on_hrs: str | None = None
    uv_intensity: str | None = None
    remark: str | None = None


class GeminiRawWaterStorageTankChlorinationEntry(BaseModel):
    time: str | None = None
    tank_number: str | None = None
    chlorine_in_ppm: str | None = None


class GeminiRawWaterFreeChlorineEntry(BaseModel):
    time: str | None = None
    chlorine: str | None = None


class GeminiChemicalDetailsEntry(BaseModel):
    particulars: str | None = None
    date: str | None = None
    kg: str | None = None
    lot_number: str | None = None
    supplier: str | None = None


class GeminiMicronFilterEntry(BaseModel):
    five_micron_differential_pressure: str | None = None
    five_micron_inlet_pressure: str | None = None
    five_micron_outlet_pressure: str | None = None
    three_micron_differential_pressure: str | None = None
    three_micron_initial_pressure: str | None = None
    three_micron_final_pressure: str | None = None
    line_details: str | None = None
    sr_no: str | None = None
    remarks: str | None = None


class GeminiBackwashingFilterEntry(BaseModel):
    stream: str | None = None

    acf_backwash_start_time: str | None = None
    acf_backwash_stop_time: str | None = None
    acf_rinsing_start_time: str | None = None
    acf_rinsing_stop_time: str | None = None
    acf_turbidity: str | None = None

    lead_backwash_start_time: str | None = None
    lead_backwash_stop_time: str | None = None
    lead_rinsing_start_time: str | None = None
    lead_rinsing_stop_time: str | None = None
    lead_turbidity: str | None = None

    mgf_psf_backwash_start_time: str | None = None
    mgf_psf_backwash_stop_time: str | None = None
    mgf_psf_rinsing_start_time: str | None = None
    mgf_psf_rinsing_stop_time: str | None = None
    mgf_psf_turbidity: str | None = None

    lag_backwash_start_time: str | None = None
    lag_backwash_stop_time: str | None = None
    lag_rinsing_start_time: str | None = None
    lag_rinsing_stop_time: str | None = None
    lag_turbidity: str | None = None

    softener_backwash_start_time: str | None = None
    softener_backwash_stop_time: str | None = None
    softener_regeneration_start_time: str | None = None
    softener_regeneration_stop_time: str | None = None
    softener_th: str | None = None
    remarks: str | None = None


class GeminiWTPRightPageOCR(BaseModel):
    date: str | None = None
    raw_water_analysis_entries: list[GeminiRawWaterAnalysisEntry] = Field(default_factory=list)
    raw_water_analysis_column_frequencies: list[GeminiColumnFrequency] = Field(default_factory=list)
    soft_water_analysis_entries: list[GeminiSoftWaterAnalysisEntry] = Field(default_factory=list)
    soft_water_analysis_column_frequencies: list[GeminiColumnFrequency] = Field(default_factory=list)
    treated_water_analysis_entries: list[GeminiTreatedWaterAnalysisEntry] = Field(default_factory=list)
    treated_water_analysis_column_frequencies: list[GeminiColumnFrequency] = Field(default_factory=list)
    uv_light_entries: list[GeminiUVLightEntry] = Field(default_factory=list)
    raw_water_storage_tank_chlorination_entries: list[GeminiRawWaterStorageTankChlorinationEntry] = Field(
        default_factory=list
    )
    raw_water_free_chlorine_entries: list[GeminiRawWaterFreeChlorineEntry] = Field(default_factory=list)
    raw_water_free_chlorine_frequency: str | None = None
    chemical_details_entries: list[GeminiChemicalDetailsEntry] = Field(default_factory=list)
    micron_filter_entries: list[GeminiMicronFilterEntry] = Field(default_factory=list)
    backwashing_filter_entries: list[GeminiBackwashingFilterEntry] = Field(default_factory=list)


WTP_RIGHT_PAGE_PROMPT = """
You are transcribing a single photographed page from a handwritten Water Treatment Plant (WTP) daily
monitoring logbook, titled "WTP Daily Monitoring Record". The page contains one date field and 9
separate handwritten data tables. 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 tables are arranged in three vertical bands, each containing multiple
tables side by side, plus one full-width band at the bottom:
- Top band: Date box (far left) | RAW WATER ANALYSIS (wide) | SOFT WATER ANALYSIS (wide, right of
  raw water) | Raw Water Free Chlorine (far right edge, narrow, just 2 columns: Time / value).
- Middle band: TREATED WATER ANALYSIS (wide, left) | UV Light (right of treated water) | Raw Water
  Storage Tank Chlorination (narrow, right of UV Light, left of a "Remarks" shift-names box).
- Bottom band: CHEMICAL DETAILS (wide, left) | Micron Filters (right).
- Full-width band at the very bottom of the page: Backwashing of Filters / Softener.

Two tables — Raw Water Free Chlorine and Raw Water Storage Tank Chlorination — are narrow, two/three
column (Time | value, or Time | Tank No. | value) tables that run for many more rows than the wider
table(s) beside them, so they visually spill past where their neighbor ends. Raw Water Free Chlorine
in particular runs down the full right edge of the page, alongside the Soft Water Analysis, Treated
Water Analysis, and UV Light bands in turn. Treat each of these narrow tables as its own independent
list, column-bound only by its own vertical rule lines: never borrow a time or value from whatever
wide table happens to be alongside it on the same horizontal line, and never stop transcribing one
of these narrow tables just because a neighboring wider table has ended above or beside it.

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. raw_water_analysis_entries: rows of the "RAW WATER ANALYSIS" table (top-left). Columns, left to
   right: time, raw_water_storage_tank_use, odour, appearance, m_alk, th, cah, pH, tds,
   total_chlorine, p_alk, turbidity, sulphate, chloride, iron.
   Position: top-left, first table on the page.
   Expected rows: most columns are specced "Daily" — normally just 1 filled row per day. Actual
   count is very often 0 or 1; do not pad toward a higher count.
   raw_water_analysis_column_frequencies: this table prints its own "Frequency" column showing
   each row-parameter's cadence (most read "Daily", but a few, e.g. Sulphate/Chloride/Iron, are
   commonly "Weekly" — it varies by form, don't assume). For each column above whose printed
   frequency label is legible, add one {"column": <field name>, "frequency": <printed label>}
   entry to this list (e.g. {"column": "sulphate", "frequency": "Weekly"}). Skip a column
   entirely if no frequency is printed or it's illegible — do not guess or default it to match
   the other columns. This is printed boilerplate, not a handwritten row value, so it doesn't
   count toward "Expected rows".

3. soft_water_analysis_entries: rows of the "SOFT WATER ANALYSIS" table (top-right, next to raw
   water analysis). Columns, left to right: time, taste_odour_appearance, th, cah,
   residual_chlorine, tds, pH, turbidity, m_alk, iron, chloride.
   Position: top band, right of raw_water_analysis.
   Expected rows: dominant cadence is "Every 4 hour" — up to ~6 rows/day (roughly one every 4
   hours around the clock). Fewer filled rows is common and fine.
   soft_water_analysis_column_frequencies: same idea as raw_water_analysis_column_frequencies
   above — one {"column": <field name>, "frequency": <printed label>} entry per column whose
   frequency is legible (this table commonly reads "Every 4 hour" for most columns but "Every 8
   hour" for taste_odour_appearance; don't assume, transcribe what's actually printed per
   column, and skip a column if nothing is printed/legible for it).

4. raw_water_free_chlorine_entries: rows of the separate, narrow "Raw Water Free Chlorine" table at
   the far top-right (header "Frequency: Hourly", spec "3-5ppm"). Columns: time, chlorine.
   Position: far top-right corner, a narrow 2-column (Time, value) table that is visibly taller
   than the tables beside it and keeps running down the right edge of the page alongside the
   Soft Water Analysis, Treated Water Analysis, and UV Light bands below — see "Page layout
   overview" above. It is its own independent list.
   Expected rows: cadence is Hourly — up to ~24 rows/day, but commonly fewer (a shift may not
   run the full 24 hours, or logging may stop partway down). Do not borrow a time from any
   neighboring table for this list, and do not truncate it early just because a table beside it
   (e.g. UV Light) ends first.
   raw_water_free_chlorine_frequency: the printed cadence label in this table's own header (e.g.
   "Frequency: Hourly" -> "Hourly"), transcribed once for the whole table — this is printed
   boilerplate, not a handwritten row, so it does not count toward "Expected rows" above and is
   unrelated to the STD SPEC/other printed text ignored elsewhere on this page. null if the
   table has no printed frequency label or it's illegible.

5. treated_water_analysis_entries: rows of the "TREATED WATER ANALYSIS" table. Columns, left to
   right: time, appearance, odour, taste, p_alkalinity, m_alkalinity, total_chlorine, free_chlorine,
   pH, turbidity, tds, th, cah, aluminium, iron, chloride, sulphate,
   total_combined_sulphate_and_chloride, remark.
   Position: middle band, left, below Raw/Soft Water Analysis.
   Expected rows: dominant cadence is "4 Hourly" — up to ~6 rows/day. Fewer is common.
   treated_water_analysis_column_frequencies: same idea as raw_water_analysis_column_frequencies
   above — one {"column": <field name>, "frequency": <printed label>} entry per column whose
   frequency is legible (this table commonly reads "4 Hourly" for most columns but "Daily" for a
   few, e.g. TH/Aluminium/CAH/Iron/Chloride/Sulfate; don't assume, transcribe what's actually
   printed per column, and skip a column if nothing is printed/legible for it).

6. uv_light_entries: rows of the "UV Light" table. Columns, left to right: time,
   uv_light_number ("UV Light No."), uv_intensity ("UV Intensity"), uv_light_on_hrs
   ("UV Lamp on Hrs."), remark.
   Position: middle band, right of treated_water_analysis, left of Raw Water Storage Tank
   Chlorination.
   Visual quirk: rows are grouped by UV Light No. Within a group, only the FIRST row has time,
   uv_light_number, and uv_intensity filled in; the remaining rows of that same group normally
   leave those three columns blank and only carry an incrementing uv_light_on_hrs meter reading.
   That is expected — do not copy the first row's time/light-number/intensity down into the
   blank-looking rows, and do not mistake a lone number in the "UV Lamp on Hrs." column for
   belonging to a different table.
   Expected rows: typically 2 UV light numbers with up to ~4 on-hours readings logged per light
   per day — up to ~8 rows total, grouped consecutively by light number. Fewer is common.

7. raw_water_storage_tank_chlorination_entries: rows of the "Raw Water Storage Tank Chlorination"
   table (next to UV Light). Columns: time, tank_number ("Tank No."), chlorine_in_ppm
   ("Chlorine in PPM").
   Position: narrow column immediately right of UV Light, left of a "Remarks" shift-names box.
   Visual quirk: like UV Light, the time column is usually only filled on the first row of a
   shift check — a blank time cell does not mean the row belongs to a different table.
   Expected rows: one row per raw water storage tank at the site, numbered sequentially starting
   at 01 (e.g. 01, 02, 03, ...). The number of tanks varies by plant (this form's example runs
   01-09); treat the actual filled count as the ground truth even if lower.

8. chemical_details_entries: rows of the "CHEMICAL DETAILS" table. Each physical row is one chemical
   (e.g. Calcium Hypochlorite, Common Salt, MDC 220, Indion 8111, Indion 3100, Indion 5010, Citric
   Acid, Soda Ash, etc.) and has FIVE repeated groups of (Date, Kg, B.No./Lot No.) columns, plus one
   shared Supplier column at the end. Emit ONE entry per FILLED (Date, Kg, B.No.) group — e.g. if a
   row has 2 of the 5 groups filled in, emit 2 entries, each with that row's particulars and supplier
   but that group's own date/kg/lot_number. Skip groups that are entirely blank. Fields per entry:
   particulars, date, kg, lot_number, supplier.
   Position: bottom band, left.
   Expected rows: no fixed target — each chemical row can have 0-5 filled groups; count only what
   is actually written.

9. micron_filter_entries: rows of the "Micron Filters" table. Columns, left to right: sr_no
   ("Sr. No."), line_details ("Line Details"), then for "5 Micron": five_micron_inlet_pressure,
   five_micron_outlet_pressure, five_micron_differential_pressure; then for "3 Micron":
   three_micron_initial_pressure (labelled "Inlet Pressure"), three_micron_final_pressure
   (labelled "Outlet Pressure"), three_micron_differential_pressure; then remarks.
   Position: bottom band, right of Chemical Details.
   Expected rows: one per filter line (Sr. No.) in use; commonly just 1-3 rows, varies by plant.

10. backwashing_filter_entries: rows of the "Backwashing of Filters" table. Columns: stream
    ("Stream No."), then for each of MGF/PSF, ACF, LEAD, LAG (prefix mgf_psf, acf, lead, lag
    respectively): <prefix>_backwash_start_time, <prefix>_backwash_stop_time,
    <prefix>_rinsing_start_time, <prefix>_rinsing_stop_time, <prefix>_turbidity; then for Softener
    (no turbidity column): softener_backwash_start_time, softener_backwash_stop_time,
    softener_regeneration_start_time, softener_regeneration_stop_time, softener_th; then remarks.
    Position: full-width band at the very bottom of the page, below Chemical Details/Micron
    Filters.
    Expected rows: one per Stream No. listed; varies by plant (commonly a handful of streams).

General rules for every table:
- Transcribe each handwritten cell's text EXACTLY as written (same digits, abbreviations,
  punctuation). 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, or crossed out with a diagonal line indicating no
  further entries for that table on this page.
- Only emit as many rows per table as are actually filled in; do not invent or pad rows. Every
  "Expected rows" figure above is a ceiling based on the printed frequency/spec, not a target — the
  actual filled count on a given day is very often lower, and that's normal, not an error.
- 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_right_page_ocr_gemini(
    file_path: str,
    mime_type: str = "image/jpeg",
    model: str = DEFAULT_GEMINI_MODEL,
    raw_gemini_output_path: str | None = None,
) -> OCRResult[WTPRightPageOCR]:
    start = time_module.perf_counter()

    raw_result = generate_structured_content(
        file_path=file_path,
        mime_type=mime_type,
        prompt=WTP_RIGHT_PAGE_PROMPT,
        response_schema=GeminiWTPRightPageOCR,
        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))

    try:
        result = convert_model(WTPRightPageParser(), raw_result, WTPRightPageOCR)
    except ValidationError as exc:
        raise OcrError(
            OcrFailureCode.EXTRACTION_VALIDATION_FAILED,
            f"Extracted data failed validation: {exc}",
        ) from exc

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