from decimal import Decimal
from datetime import time
from pydantic import BaseModel, ConfigDict, Field
from typing import Annotated

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

class RawWaterAnalysisEntry(TimedLogEntry):
    raw_water_storage_tank_use: OcrField[Decimal] | None = None
    appearance: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["Nm"]), ocr_default="Nm")] = None
    odour: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["Nm"]), ocr_default="Nm")] = None
    cah: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=200))] = None
    chloride: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=250))] = None
    iron: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=0.1))] = None
    m_alk: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=200))] = None
    pH: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=6.5, max=8.5))] = None
    p_alk: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=2, max=5))] = None
    sulphate: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=250))] = None
    tds: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=500))] = None
    th: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=500))] = None
    turbidity: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=5))] = None
    total_chlorine: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=10))] = None

class SoftWaterAnalysisEntry(TimedLogEntry):
    cah: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=10))] = None
    m_alk: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=85))] = None
    residual_chlorine: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=1, max=3))] = None
    taste_odour_appearance: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["Noo/Nm"]), ocr_default="Noo/Nm")] = None
    tds: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=500))] = None
    th: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=10))] = None
    turbidity: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=0.3))] = None
    pH: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=6.5, max=8.5))] = None
    iron: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=0.3))] = None
    chloride: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=250))] = None

class TreatedWaterAnalysisEntry(TimedLogEntry):
    aluminium: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=0.2))] = None
    appearance: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["NVC"]), ocr_default="NVC")] = None
    cah: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=200))] = None
    free_chlorine: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=0))] = None
    iron: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=0.1))] = None
    m_alkalinity: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=85))] = None
    odour: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["NOO"]), ocr_default="NOO")] = None
    pH: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=6.5, max=8.5))] = None
    p_alkalinity: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=5))] = None
    taste: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["NOT"]), ocr_default="NOT")] = None
    tds: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=500))] = None
    th: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=250))] = None
    total_chlorine: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=0))] = None
    turbidity: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=0.3))] = None
    chloride: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=250))] = None
    sulphate: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=250))] = None
    total_combined_sulphate_and_chloride: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=400))] = None
    remark: OcrField[str] | None = None

class UVLightEntry(TimedLogEntry):
    uv_light_number: OcrField[Decimal] | None = None
    uv_light_on_hrs: OcrField[Decimal] | None = None
    uv_intensity: OcrField[Decimal] | None = None
    remark: OcrField[str] | None = None


class RawWaterStorageTankChlorinationEntry(TimedLogEntry):
    tank_number: OcrField[Decimal] | None = None
    chlorine_in_ppm: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=5))] = None

class RawWaterFreeChlorineEntry(TimedLogEntry):
    chlorine: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=3, max=5))] = None

class ChemicalDetailsEntry(SourceTaggedEntry):
    particulars: OcrField[str] | None = None
    date: OcrField[str] | None = None
    kg: OcrField[str] | None = None
    lot_number: OcrField[str] | None = None
    supplier: OcrField[str] | None = None


class MicronFilterEntry(SourceTaggedEntry):
    five_micron_differential_pressure: OcrField[Decimal] | None = None
    five_micron_inlet_pressure: OcrField[Decimal] | None = None
    five_micron_outlet_pressure: OcrField[Decimal] | None = None

    three_micron_differential_pressure: OcrField[Decimal] | None = None
    three_micron_initial_pressure: OcrField[Decimal] | None = None
    three_micron_final_pressure: OcrField[Decimal] | None = None

    line_details: OcrField[Decimal] | None = None
    sr_no: OcrField[Decimal] | None = None
    remarks: OcrField[str] | None = None


class BackwashingFilterEntry(SourceTaggedEntry):
    # Only `stream` is required -- a row's five treatment groups (ACF/Lead/MGF-PSF/Lag/
    # Softener) legitimately vary per stream, so a blank group isn't necessarily an OCR
    # miss the way a blank `stream` is. The logbook also prints no spec/range for any of
    # this table's values (unlike e.g. Free Chlorine's "3-5ppm"), so there's no basis for
    # a RangeSpec on the turbidity/softener_th fields either -- required=True on `stream`
    # is what's left to at least flag when DocAI drops a row's identifying column.
    stream: Annotated[OcrField[str] | None, FieldSpec(required=True)] = None

    acf_backwash_start_time: OcrField[time] | None = None
    acf_backwash_stop_time: OcrField[time] | None = None
    acf_rinsing_start_time: OcrField[time] | None = None
    acf_rinsing_stop_time: OcrField[time] | None = None
    acf_turbidity: OcrField[Decimal] | None = None

    lead_backwash_start_time: OcrField[time] | None = None
    lead_backwash_stop_time: OcrField[time] | None = None
    lead_rinsing_start_time: OcrField[time] | None = None
    lead_rinsing_stop_time: OcrField[time] | None = None
    lead_turbidity: OcrField[Decimal] | None = None

    mgf_psf_backwash_start_time: OcrField[time] | None = None
    mgf_psf_backwash_stop_time: OcrField[time] | None = None
    mgf_psf_rinsing_start_time: OcrField[time] | None = None
    mgf_psf_rinsing_stop_time: OcrField[time] | None = None
    mgf_psf_turbidity: OcrField[Decimal] | None = None

    lag_backwash_start_time: OcrField[time] | None = None
    lag_backwash_stop_time: OcrField[time] | None = None
    lag_rinsing_start_time: OcrField[time] | None = None
    lag_rinsing_stop_time: OcrField[time] | None = None
    lag_turbidity: OcrField[Decimal] | None = None

    softener_backwash_start_time: OcrField[time] | None = None
    softener_backwash_stop_time: OcrField[time] | None = None
    softener_regeneration_start_time: OcrField[time] | None = None
    softener_regeneration_stop_time: OcrField[time] | None = None
    softener_th: OcrField[Decimal] | None = None
    remarks: OcrField[str] | None = None

class WTPRightPageOCR(LogbookOCR):
    raw_water_analysis_entries: list[RawWaterAnalysisEntry] = Field(default_factory=list)
    # Printed per-column "Frequency" label for each column in raw_water_analysis_entries,
    # keyed by RawWaterAnalysisEntry field name (e.g. {"sulphate": "Weekly"}) -- frequency is
    # usually the same across a table's columns but not always (this table mixes Daily and
    # Weekly), so it's transcribed per column rather than assumed. Omit a key when that
    # column's frequency isn't printed/legible. Gemini-only -- Document AI has no
    # frequency-related entity/property anywhere on this page (confirmed against
    # test_data/wtp_right_page's raw entities), so this stays None on the DocAI path.
    raw_water_analysis_column_frequencies: dict[str, str] | None = None
    soft_water_analysis_entries: list[SoftWaterAnalysisEntry] = Field(default_factory=list)
    soft_water_analysis_column_frequencies: dict[str, str] | None = None
    treated_water_analysis_entries: list[TreatedWaterAnalysisEntry] = Field(default_factory=list)
    treated_water_analysis_column_frequencies: dict[str, str] | None = None
    uv_light_entries: list[UVLightEntry] = Field(default_factory=list)
    raw_water_storage_tank_chlorination_entries: list[RawWaterStorageTankChlorinationEntry] = Field(default_factory=list)
    raw_water_free_chlorine_entries: list[RawWaterFreeChlorineEntry] = Field(default_factory=list)
    # The one table on this page whose printed "Frequency: ..." label is genuinely
    # single-column (unlike e.g. Raw/Soft Water Analysis, where FREQUENCY is printed
    # per-column and doesn't reduce to one label for the whole table).
    # Gemini-only for now: the Document AI processor for this page has no entity/property
    # type for this label at all (confirmed against test_data/wtp_right_page's raw
    # entities), so it stays None on the DocAI path until the processor is retrained.
    raw_water_free_chlorine_frequency: OcrField[str] | None = None
    chemical_details_entries: list[ChemicalDetailsEntry] = Field(default_factory=list)
    micron_filter_entries: list[MicronFilterEntry] = Field(default_factory=list)
    backwashing_filter_entries: list[BackwashingFilterEntry] = Field(default_factory=list)
    model_config = ConfigDict(populate_by_name=True)

class WTPRightPageParser(OcrParser):
    _ENTRY_TYPES: dict[str, tuple[str, type[BaseModel]]] = {
        "raw_water_analysis_entries": (
            "raw_water_analysis_entries",
            RawWaterAnalysisEntry,
        ),
        "soft_water_analysis_entries": (
            "soft_water_analysis_entries",
            SoftWaterAnalysisEntry,
        ),
        "treated_water_analysis_entries": (
            "treated_water_analysis_entries",
            TreatedWaterAnalysisEntry,
        ),
        "uv_light_entries": (
            "uv_light_entries",
            UVLightEntry,
        ),
        "raw_water_storage_tank_chlorination_entries": (
            "raw_water_storage_tank_chlorination_entries",
            RawWaterStorageTankChlorinationEntry,
        ),
        "raw_water_free_chlorine_entries": (
            "raw_water_free_chlorine_entries",
            RawWaterFreeChlorineEntry,
        ),
        "chemical_detail_entries": (
            "chemical_details_entries",
            ChemicalDetailsEntry,
        ),
        "micro_filter_entries": (
            "micron_filter_entries",
            MicronFilterEntry,
        ),
        "backwashing_filter_entries": (
            "backwashing_filter_entries",
            BackwashingFilterEntry,
        ),
    }

    def parse(self, raw_ocr_document: RawOcrDocument) -> WTPRightPageOCR:
        parsed_document = WTPRightPageOCR(
            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