import time as time_module

from pydantic import BaseModel, Field, ValidationError

from ..exceptions import OcrError, OcrFailureCode
from ..models import OCRResult
from ..parse_ocr.wtp_coagulation_parser import WTPCoagulationLogOCR, WTPCoagulationLogParser
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 GeminiRawWaterTankEntry(BaseModel):
    time: str | None = None
    odour: str | None = None
    appearance: str | None = None
    turbidity: str | None = None
    pH: str | None = None
    p_value: str | None = None
    m_value: str | None = None
    th: str | None = None
    cah: str | None = None
    total_chlorine: str | None = None
    tds: str | None = None
    toc: str | None = None
    iron: str | None = None
    chloride: str | None = None
    sulphate: str | None = None


class GeminiAfterPSFEntry(BaseModel):
    time: str | None = None
    odour: str | None = None
    appearance: str | None = None
    turbidity_psf_1: str | None = None
    turbidity_psf_2: str | None = None
    pH: str | None = None
    p_value: str | None = None
    m_value: str | None = None
    th: str | None = None
    cah: str | None = None
    free_cl2: str | None = None


class GeminiAfterACFEntry(BaseModel):
    time: str | None = None
    taste: str | None = None
    odour: str | None = None
    appearance: str | None = None
    turbidity: str | None = None
    pH: str | None = None
    tds: str | None = None
    m_value: str | None = None
    th: str | None = None
    cah: str | None = None
    total_cl2: str | None = None


class GeminiRawWaterChlorineEntry(BaseModel):
    time: str | None = None
    free_chlorine: str | None = None


class GeminiAfter5MicronEntry(BaseModel):
    time: str | None = None
    taste: str | None = None
    odour: str | None = None
    appearance: str | None = None
    turbidity: str | None = None
    pH: str | None = None
    tds: str | None = None
    m_value: str | None = None
    th: str | None = None
    cah: str | None = None
    free_cl2: str | None = None
    total_cl2: str | None = None
    iron: str | None = None
    chloride: str | None = None
    sulphate: str | None = None
    combined_sulphate_and_chloride: str | None = None


class GeminiBackwashEntry(BaseModel):
    filter_name: str | None = None
    frequency: str | None = None
    pressure_before_bw: str | None = None
    turbidity_before_bw: str | None = None
    backwash_start_time: str | None = None
    backwash_end_time: str | None = None
    rinsewash_start_time: str | None = None
    rinsewash_end_time: str | None = None
    pressure_after_bw: str | None = None
    turbidity_after_bw: str | None = None


class GeminiChemicalsPreparationDetailEntry(BaseModel):
    particulars: str | None = None
    time: str | None = None
    batch_or_lot_number: str | None = None
    supplier: str | None = None
    kg: str | None = None


class GeminiChemicalDosingEntry(BaseModel):
    time: str | None = None
    lime: str | None = None
    feso4: str | None = None
    cacl2: str | None = None
    soda_ash: str | None = None


class GeminiWTPCoagulationLogOCR(BaseModel):
    date: str | None = None
    raw_water_tank_entries: list[GeminiRawWaterTankEntry] = Field(default_factory=list)
    raw_water_tank_column_frequencies: list[GeminiColumnFrequency] = Field(default_factory=list)
    after_psf_entries: list[GeminiAfterPSFEntry] = Field(default_factory=list)
    after_psf_column_frequencies: list[GeminiColumnFrequency] = Field(default_factory=list)
    after_acf_entries: list[GeminiAfterACFEntry] = Field(default_factory=list)
    after_acf_column_frequencies: list[GeminiColumnFrequency] = Field(default_factory=list)
    raw_water_chlorine_entries: list[GeminiRawWaterChlorineEntry] = Field(default_factory=list)
    raw_water_chlorine_frequency: str | None = None
    after_5_micron_entries: list[GeminiAfter5MicronEntry] = Field(default_factory=list)
    after_5_micron_column_frequencies: list[GeminiColumnFrequency] = Field(default_factory=list)
    backwash_entries: list[GeminiBackwashEntry] = Field(default_factory=list)
    chemicals_preparation_detail_entries: list[GeminiChemicalsPreparationDetailEntry] = Field(
        default_factory=list
    )
    chemical_dosing_entries: list[GeminiChemicalDosingEntry] = Field(default_factory=list)


WTP_COAGULATION_PROMPT = """
You are transcribing a single photographed page from a handwritten Water Treatment Plant (WTP)
"Treated Water Monitoring Record - Coagulation System" logbook ("Enrich Agro Food Products Pvt. Ltd.
Rohtak", form QSE-8.1-FRM-WTP-1.1.1.20). The page contains one date field and 8 handwritten data
tables/sections. Your ONLY job is literal, verbatim transcription of what is written in each cell —
do not calculate, round, unit-convert, or "correct" any value. Ignore all printed "STD SPEC" rows —
only transcribe the handwritten "Time"/"Result" rows actually filled in. The printed "FREQUENCY"
row is also boilerplate, not a handwritten row, but for the tables noted below it must still be
transcribed into that table's own *_column_frequencies field (never as a handwritten Time/Result
entry) — see each numbered item.

Page layout overview: the tables run top-to-bottom down the page in this order: RAW WATER TANK-1
(top of page) → AFTER PSF, with the narrow RAW WATER CHLORINE / FREE CL2 table at the far right
of that same row → AFTER ACF → AFTER 5 MICRON → BACKWASH → CHEMICALS PREPARATION DETAIL (near the
backwash table; may appear twice, e.g. for two shifts) → CHEMICAL DOSING. Each table/section is
its own independent list — never borrow a time or value from a neighboring table just because it
sits nearby on the page, and don't stop transcribing one table because a different one above or
below it has ended.

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

1. date: the handwritten date near the top-left of the page, exactly as written (e.g. "19/5/26").
   null if blank or illegible.

2. raw_water_tank_entries: row(s) of the "RAW WATER TANK-1" table (top of page). Columns, left to
   right: time, odour, appearance, turbidity, pH, p_value ("P Value"), m_value ("M-Value"), th
   ("TH"), cah ("CaH"), total_chlorine, tds, toc ("TOC"), iron, chloride, sulphate.
   raw_water_tank_column_frequencies: this table prints its own "FREQUENCY" row above the
   "Result" rows, giving each column's own cadence (mostly "Daily", but a column or two, e.g.
   sulphate, may read "WEEKLY" — it varies by form, don't assume). For each column above whose
   printed frequency 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 nothing is printed/legible for it — never guess or default it to match the other
   columns.

3. after_psf_entries: rows of the "AFTER PSF" table. Columns, left to right: time, odour,
   appearance, turbidity_psf_1 ("Turbidity PSF-1"), turbidity_psf_2 ("Turbidity PSF-2"), pH,
   p_value ("P Value"), m_value ("M-Value"), th ("TH"), cah ("CaH"), free_cl2 ("Free CL2").
   after_psf_column_frequencies: same idea as raw_water_tank_column_frequencies above — one
   {"column": <field name>, "frequency": <printed label>} entry per column whose "FREQUENCY" row
   value is legible. Skip a column if nothing is printed/legible for it.

4. raw_water_chlorine_entries: rows of the narrow "RAW WATER CHLORINE / FREE CL2" table at the
   far right of the "AFTER PSF" row (multiple Time/Result column-pairs, hourly readings). Columns:
   time, free_chlorine ("Result").
   raw_water_chlorine_frequency: unlike the wider tables on this page, this table is a single
   column (just "Result"), so its printed "FREQUENCY" label (e.g. "Hourly") is one value for the
   whole table, not a per-column dict. Transcribe it here once. null if not printed/legible.

5. after_acf_entries: rows of the "AFTER ACF" table. Columns, left to right: time, taste, odour,
   appearance, turbidity, pH, tds ("TDS"), m_value ("M-Value"), th ("TH"), cah ("CaH"), total_cl2
   ("Total CL2").
   after_acf_column_frequencies: same idea as raw_water_tank_column_frequencies above — one
   {"column": <field name>, "frequency": <printed label>} entry per column whose "FREQUENCY" row
   value is legible. Skip a column if nothing is printed/legible for it.

6. after_5_micron_entries: rows of the "AFTER 5 MICRON" table. Columns, left to right: time, taste,
   odour, appearance, turbidity, pH, tds ("TDS"), m_value ("M-Value"), th ("TH"), cah ("CaH"),
   free_cl2 ("FREE CL2"), total_cl2 ("Total CL2"), iron ("IRON"), chloride ("CHLORIDE"), sulphate
   ("SULPHATE"), combined_sulphate_and_chloride ("Combined Sulphate & Chloride").
   after_5_micron_column_frequencies: same idea as raw_water_tank_column_frequencies above —
   this table commonly reads "4 HRS" for most columns but "Daily" for a few (e.g. iron/chloride/
   sulphate/combined_sulphate_and_chloride). One {"column": <field name>, "frequency": <printed
   label>} entry per column whose value is legible. Skip a column if nothing is printed/legible
   for it.

7. backwash_entries: rows of the "BACKWASH" table, one entry per filter row label (e.g. "PSF1",
   "PSF2", "ACF", "LEAD ACF", "LAG ACF"). Fields: filter_name (the row label), frequency (this
   table's printed "FREQUENCY" column is per FILTER ROW, not per metric column — e.g. "24 HRS"
   for the PSF filters, "96 HRS" for the ACF filters — so transcribe it as this row's own value,
   not a table-level dict), pressure_before_bw
   ("Pressure Before BW"), turbidity_before_bw ("Turbidity Before BW"), backwash_start_time
   ("Backwash Start"), backwash_end_time ("Backwash End"), rinsewash_start_time ("Rinsewash Start"),
   rinsewash_end_time ("Rinsewash End"), pressure_after_bw ("Pressure After BW"),
   turbidity_after_bw ("Turbidity After BW"). Only emit a row if at least one of its cells is
   filled in.

8. chemicals_preparation_detail_entries: rows of the "CHEMICALS PREPARATION DETAIL" table(s) (near
   the backwash table), one entry per filled-in particular row (e.g. "Ferrous Sulfate", "Lime",
   "Cal. Hypochlorite", "Soda Ash"). Fields: particulars, time, kg ("Kg"), batch_or_lot_number
   ("B.No./Lot No."), supplier. If this section appears twice on the page (e.g. two shifts), emit
   one entry per filled row across both occurrences.

9. chemical_dosing_entries: rows of the "CHEMICAL DOSING" table. Columns, left to right: time,
   lime ("Lime (2%) LPH"), feso4 ("FeSO4 (6%) LPH"), cacl2 ("CaCl2 (3%) LPH"), soda_ash
   ("Soda ash (3%) LPH").

Ignore the "EQUIPMENT RUNNING" section (Coagulation Start/End Time, Purge Removal, Shift, Operator/
Executive Name, signatures) and any printed boilerplate (form code, revision/issue date) — 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. "Noo" or "NVC"). 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 (e.g. a dash "—" or "-" placeholder
  row with no further readings).
- 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_coagulation_log_ocr_gemini(
    file_path: str,
    mime_type: str = "image/jpeg",
    model: str = DEFAULT_GEMINI_MODEL,
    raw_gemini_output_path: str | None = None,
) -> OCRResult[WTPCoagulationLogOCR]:
    start = time_module.perf_counter()

    raw_result = generate_structured_content(
        file_path=file_path,
        mime_type=mime_type,
        prompt=WTP_COAGULATION_PROMPT,
        response_schema=GeminiWTPCoagulationLogOCR,
        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(WTPCoagulationLogParser(), raw_result, WTPCoagulationLogOCR)
    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,
    )
