from decimal import Decimal
from datetime import time as Time
from typing import Annotated

from pydantic import BaseModel, ConfigDict, Field

from ..models import FieldSpec, OcrField, OptionSpec, RawOcrDocument
from .parser import SourceTaggedEntry, TimedLogEntry, LogbookOCR, OcrParser


class RawWaterTankEntry(TimedLogEntry):
    odour: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["NOO"]), ocr_default="NOO")] = None
    appearance: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["NVC"]), ocr_default="NVC")] = None
    turbidity: OcrField[Decimal] | None = None
    pH: OcrField[Decimal] | None = None
    p_value: OcrField[Decimal] | None = None
    m_value: OcrField[Decimal] | None = None
    th: OcrField[Decimal] | None = None
    cah: OcrField[Decimal] | None = None
    total_chlorine: OcrField[Decimal] | None = None
    tds: OcrField[Decimal] | None = None
    toc: OcrField[Decimal] | None = None
    iron: OcrField[Decimal] | None = None
    chloride: OcrField[Decimal] | None = None
    sulphate: OcrField[Decimal] | None = None


class AfterPSFEntry(TimedLogEntry):
    odour: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["NOO"]), ocr_default="NOO")] = None
    appearance: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["NVC"]), ocr_default="NVC")] = None
    turbidity_psf_1: OcrField[Decimal] | None = None
    turbidity_psf_2: OcrField[Decimal] | None = None
    pH: OcrField[Decimal] | None = None
    p_value: OcrField[Decimal] | None = None
    m_value: OcrField[Decimal] | None = None
    th: OcrField[Decimal] | None = None
    cah: OcrField[Decimal] | None = None
    free_cl2: OcrField[Decimal] | None = None


class AfterACFEntry(TimedLogEntry):
    taste: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["NOT"]), ocr_default="NOT")] = None
    odour: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["NOO"]), ocr_default="NOO")] = None
    appearance: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["NVC"]), ocr_default="NVC")] = None
    turbidity: OcrField[Decimal] | None = None
    pH: OcrField[Decimal] | None = None
    tds: OcrField[Decimal] | None = None
    m_value: OcrField[Decimal] | None = None
    th: OcrField[Decimal] | None = None
    cah: OcrField[Decimal] | None = None
    total_cl2: OcrField[Decimal] | None = None


class RawWaterChlorineEntry(TimedLogEntry):
    free_chlorine: OcrField[Decimal] | None = None


class After5MicronEntry(TimedLogEntry):
    taste: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["NOT"]), ocr_default="NOT")] = None
    odour: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["NOO"]), ocr_default="NOO")] = None
    appearance: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["NVC"]), ocr_default="NVC")] = None
    turbidity: OcrField[Decimal] | None = None
    pH: OcrField[Decimal] | None = None
    tds: OcrField[Decimal] | None = None
    m_value: OcrField[Decimal] | None = None
    th: OcrField[Decimal] | None = None
    cah: OcrField[Decimal] | None = None
    free_cl2: OcrField[Decimal] | None = None
    total_cl2: OcrField[Decimal] | None = None

    iron: OcrField[Decimal] | None = None
    chloride: OcrField[Decimal] | None = None
    sulphate: OcrField[Decimal] | None = None
    combined_sulphate_and_chloride: OcrField[Decimal] | None = None


class BackwashEntry(SourceTaggedEntry):
    filter_name: OcrField[str] | None = None
    # Unlike the analysis tables' per-column frequency, this table's printed "FREQUENCY"
    # column is per-row -- each filter (filter_name) has its own backwash cadence (e.g.
    # "24 HRS" for PSF filters, "96 HRS" for ACF filters), not a value shared down a column.
    frequency: OcrField[str] | None = None

    pressure_before_bw: OcrField[Decimal] | None = None
    turbidity_before_bw: OcrField[Decimal] | None = None

    backwash_start_time: OcrField[Time] | None = None
    backwash_end_time: OcrField[Time] | None = None

    rinsewash_start_time: OcrField[Time] | None = None
    rinsewash_end_time: OcrField[Time] | None = None

    pressure_after_bw: OcrField[Decimal] | None = None
    turbidity_after_bw: OcrField[Decimal] | None = None


class ChemicalsPreparationDetailEntry(SourceTaggedEntry):
    particulars: OcrField[str] | None = None
    time: OcrField[Time] | None = None
    batch_or_lot_number: OcrField[str] | None = None
    supplier: OcrField[str] | None = None
    kg: OcrField[Decimal] | None = None


class ChemicalDosingEntry(SourceTaggedEntry):
    time: OcrField[Time] | None = None

    lime: OcrField[Decimal] | None = None
    feso4: OcrField[Decimal] | None = None
    cacl2: OcrField[Decimal] | None = None
    soda_ash: OcrField[Decimal] | None = None


class WTPCoagulationLogOCR(LogbookOCR):
    raw_water_tank_entries: list[RawWaterTankEntry] = Field(default_factory=list)
    # Printed per-column "FREQUENCY" label for each column, keyed by the entry model's own
    # field name -- see RawWaterAnalysisEntry's equivalent field in wtp_right_page_parser.py
    # for the full rationale (mixed frequencies within one table, no standard rule for which
    # columns differ, so transcribe rather than infer). Gemini-only -- Document AI has no
    # frequency-related entity/property anywhere on this page.
    raw_water_tank_column_frequencies: dict[str, str] | None = None
    after_psf_entries: list[AfterPSFEntry] = Field(default_factory=list)
    after_psf_column_frequencies: dict[str, str] | None = None
    after_acf_entries: list[AfterACFEntry] = Field(default_factory=list)
    after_acf_column_frequencies: dict[str, str] | None = None
    raw_water_chlorine_entries: list[RawWaterChlorineEntry] = Field(default_factory=list)
    # Genuinely single-column table (only "free_chlorine"), like raw_water_free_chlorine on
    # the right page -- one label for the whole table rather than a per-column dict.
    raw_water_chlorine_frequency: OcrField[str] | None = None
    after_5_micron_entries: list[After5MicronEntry] = Field(default_factory=list)
    after_5_micron_column_frequencies: dict[str, str] | None = None
    backwash_entries: list[BackwashEntry] = Field(default_factory=list)
    chemicals_preparation_detail_entries: list[ChemicalsPreparationDetailEntry] = Field(default_factory=list)
    chemical_dosing_entries: list[ChemicalDosingEntry] = Field(default_factory=list)

    model_config = ConfigDict(populate_by_name=True)


class WTPCoagulationLogParser(OcrParser):
    _ENTRY_TYPES: dict[str, tuple[str, type[BaseModel]]] = {
        "raw_water_tank_entries": (
            "raw_water_tank_entries",
            RawWaterTankEntry,
        ),
        "after_psf_entries": (
            "after_psf_entries",
            AfterPSFEntry,
        ),
        "after_acf_entries": (
            "after_acf_entries",
            AfterACFEntry,
        ),
        "raw_water_chlorine_entries": (
            "raw_water_chlorine_entries",
            RawWaterChlorineEntry,
        ),
        "after_5_micron_entries": (
            "after_5_micron_entries",
            After5MicronEntry,
        ),
        "backwash_entries": (
            "backwash_entries",
            BackwashEntry,
        ),
        "chemicals_preparation_detail_entries": (
            "chemicals_preparation_detail_entries",
            ChemicalsPreparationDetailEntry,
        ),
        "chemical_dosing_entries": (
            "chemical_dosing_entries",
            ChemicalDosingEntry,
        ),
    }

    def parse(self, raw_ocr_document: RawOcrDocument) -> WTPCoagulationLogOCR:
        parsed_document = WTPCoagulationLogOCR(
            date=self.get_ocr_date(raw_ocr_document)
        )

        for entity in raw_ocr_document.entities:
            entry_definition = self._ENTRY_TYPES.get(entity.type)
            if entry_definition is None:
                continue

            list_field_name, entry_model = entry_definition
            parsed_entry = self.parse_entry(entity, entry_model, source_table=list_field_name)
            getattr(parsed_document, list_field_name).append(parsed_entry)

        return parsed_document