import time as time_module

from pydantic import BaseModel, Field, ValidationError

from ..exceptions import OcrError, OcrFailureCode
from ..models import OCRResult
from ..parse_ocr.etp_logbook_parser import ETPLogbookOCR, ETPLogbookParser
from .client import DEFAULT_GEMINI_MODEL, generate_structured_content
from .convert import convert_model


class GeminiProcessMonitoringEntry(BaseModel):
    time: str | None = None

    collection_tank_ph: str | None = None
    collection_tank_og: str | None = None

    eq_tank_1_ph: str | None = None
    eq_tank_1_tds: str | None = None
    eq_tank_1_bod: str | None = None
    eq_tank_1_cod: str | None = None
    eq_tank_1_tss: str | None = None

    eq_tank_2_ph: str | None = None
    eq_tank_2_tds: str | None = None
    eq_tank_2_bod: str | None = None
    eq_tank_2_cod: str | None = None
    eq_tank_2_tss: str | None = None

    aeration_tank_1_appearance: str | None = None
    aeration_tank_1_ph: str | None = None
    aeration_tank_1_do: str | None = None
    aeration_tank_1_mlss: str | None = None
    aeration_tank_1_svi: str | None = None

    aeration_tank_2_appearance: str | None = None
    aeration_tank_2_ph: str | None = None
    aeration_tank_2_do: str | None = None
    aeration_tank_2_mlss: str | None = None
    aeration_tank_2_svi: str | None = None

    aeration_tank_3_appearance: str | None = None
    aeration_tank_3_ph: str | None = None
    aeration_tank_3_do: str | None = None
    aeration_tank_3_mlss: str | None = None
    aeration_tank_3_svi: str | None = None


class GeminiTreatedWaterEntry(BaseModel):
    time: str | None = None
    appearance: str | None = None
    odour: str | None = None
    colour: str | None = None
    tds: str | None = None
    pH: str | None = None
    bod: str | None = None
    cod: str | None = None
    tss: str | None = None
    do: str | None = None
    chlorine: str | None = None
    oil_grease: str | None = None
    phosphorus: str | None = None


class GeminiETPLogbookOCR(BaseModel):
    date: str | None = None
    process_monitoring_entries: list[GeminiProcessMonitoringEntry] = Field(default_factory=list)
    treated_water_entries: list[GeminiTreatedWaterEntry] = Field(default_factory=list)


ETP_LOGBOOK_PROMPT = """
You are transcribing a single photographed page from an "ETP Daily Monitoring Record" logbook
(Effluent Treatment Plant), form code EN-321-FRM-06-01-01. The page contains one date field and
2 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. Ignore all printed
"Frequency" template rows and the printed standard-spec reference values shown under each column
header (e.g. "2500-3000 mg/ltrs", "<2000 mg/l", "Not Less than 4 mg/l") — those are fixed printed
targets, not handwritten readings, and must not be transcribed as row data.

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. "29/06/2026").
   null if blank or illegible.

2. process_monitoring_entries: one entry per handwritten "TIME" row of the top table, which spans
   the Collection Tank, Equalization Tank-1, Equalization Tank-2, Aeration Tank-1, Aeration Tank-2,
   and Aeration Tank-3 column groups. Columns, left to right:
   - collection_tank_ph ("Collection Tank" pH Value), collection_tank_og ("Collection Tank" O&G
     mg/liters)
   - eq_tank_1_ph, eq_tank_1_tds, eq_tank_1_bod, eq_tank_1_cod, eq_tank_1_tss (the "Equalization
     Tank-1" pH/TDS/BOD/COD/TSS columns)
   - eq_tank_2_ph, eq_tank_2_tds, eq_tank_2_bod, eq_tank_2_cod, eq_tank_2_tss (the "Equalization
     Tank-2" columns)
   - aeration_tank_1_appearance ("App." column, e.g. "Clear"), aeration_tank_1_ph,
     aeration_tank_1_do ("DO"), aeration_tank_1_mlss ("MLSS"), aeration_tank_1_svi ("SVI") (the
     "Aeration Tank-1" columns)
   - aeration_tank_2_appearance, aeration_tank_2_ph, aeration_tank_2_do, aeration_tank_2_mlss,
     aeration_tank_2_svi (the "Aeration Tank-2" columns)
   - aeration_tank_3_appearance, aeration_tank_3_ph, aeration_tank_3_do, aeration_tank_3_mlss,
     aeration_tank_3_svi (the "Aeration Tank-3" columns)
   Only emit a row if at least one of its cells is filled in (a row where every cell is blank or
   "-"/"—" should still be emitted as long as a TIME is written, since the time itself is data).

3. treated_water_entries: one entry per handwritten "Time" row of the second table, titled
   "Treated (Outlet) Water". Columns, left to right: time, appearance ("App.", e.g. "Clear"),
   odour ("Odur", e.g. "No off odour"), colour ("Colour"), tds ("TDS"), pH ("Ph"), bod ("BOD"),
   cod ("COD"), tss ("TSS"), do ("DO"), chlorine ("Chlorine", e.g. "NIL"), oil_grease
   ("Oil & Grease"), phosphorus ("Phosphorus").

Ignore the footer of the page (QA Executive / Environment Manager signature blocks, "Prepared By" /
"Approved BY" lines, any handwritten notes about test status like "B.O.D Result is under test", and
printed boilerplate such as classification labels, issue/revision dates) — none of that is 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. "Cloan" or "NJL"). Do not compute, round, normalize units, or infer values that
  are not written.
- Use null for any cell that is blank, illegible, or a dash "—"/"-" placeholder.
- Skip (do not emit) a row that is entirely empty, including its time column.
- 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_etp_logbook_ocr_gemini(
    file_path: str,
    mime_type: str = "image/jpeg",
    model: str = DEFAULT_GEMINI_MODEL,
    raw_gemini_output_path: str | None = None,
) -> OCRResult[ETPLogbookOCR]:
    start = time_module.perf_counter()

    raw_result = generate_structured_content(
        file_path=file_path,
        mime_type=mime_type,
        prompt=ETP_LOGBOOK_PROMPT,
        response_schema=GeminiETPLogbookOCR,
        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(ETPLogbookParser(), raw_result, ETPLogbookOCR)
    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,
    )
