from decimal import Decimal
from typing import Annotated

from pydantic import BaseModel, Field

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

ClearAppearance = Annotated[
    OcrField[str] | None,
    FieldSpec(ocr_spec=OptionSpec[str](options=["Clear"]), ocr_default="Clear"),
]


class ProcessMonitoringEntry(TimedLogEntry):
    collection_tank_ph: OcrField[Decimal] | None = None
    collection_tank_og: OcrField[Decimal] | None = None

    eq_tank_1_ph: OcrField[Decimal] | None = None
    eq_tank_1_tds: OcrField[Decimal] | None = None
    eq_tank_1_bod: OcrField[Decimal] | None = None
    eq_tank_1_cod: OcrField[Decimal] | None = None
    eq_tank_1_tss: OcrField[Decimal] | None = None

    eq_tank_2_ph: OcrField[Decimal] | None = None
    eq_tank_2_tds: OcrField[Decimal] | None = None
    eq_tank_2_bod: OcrField[Decimal] | None = None
    eq_tank_2_cod: OcrField[Decimal] | None = None
    eq_tank_2_tss: OcrField[Decimal] | None = None

    aeration_tank_1_appearance: ClearAppearance = None
    aeration_tank_1_ph: OcrField[Decimal] | None = None
    aeration_tank_1_do: OcrField[Decimal] | None = None
    aeration_tank_1_mlss: Annotated[
        OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=2500, max=3000))
    ] = None
    aeration_tank_1_svi: Annotated[
        OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=120, max=150))
    ] = None

    aeration_tank_2_appearance: ClearAppearance = None
    aeration_tank_2_ph: OcrField[Decimal] | None = None
    aeration_tank_2_do: OcrField[Decimal] | None = None
    aeration_tank_2_mlss: Annotated[
        OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=2500, max=3000))
    ] = None
    aeration_tank_2_svi: Annotated[
        OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=120, max=150))
    ] = None

    aeration_tank_3_appearance: ClearAppearance = None
    aeration_tank_3_ph: OcrField[Decimal] | None = None
    aeration_tank_3_do: OcrField[Decimal] | None = None
    aeration_tank_3_mlss: Annotated[
        OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=2500, max=3000))
    ] = None
    aeration_tank_3_svi: Annotated[
        OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=120, max=150))
    ] = None


class TreatedWaterEntry(TimedLogEntry):
    appearance: ClearAppearance = None
    odour: Annotated[
        OcrField[str] | None,
        FieldSpec(ocr_spec=OptionSpec[str](options=["No off odour"]), ocr_default="No off odour"),
    ] = None
    colour: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=100))] = None
    tds: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=2000))] = None
    pH: OcrField[Decimal] | None = None
    bod: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=30))] = None
    cod: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=250))] = None
    tss: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=50))] = None
    do: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=4, max=None))] = None
    chlorine: OcrField[str] | None = None
    oil_grease: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=10))] = None
    phosphorus: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=2))] = None


class ETPLogbookOCR(LogbookOCR):
    process_monitoring_entries: list[ProcessMonitoringEntry] = Field(default_factory=list)
    treated_water_entries: list[TreatedWaterEntry] = Field(default_factory=list)


class ETPLogbookParser(OcrParser):
    _ENTRY_TYPES: dict[str, tuple[str, type[BaseModel]]] = {
        "process_monitoring_entries": (
            "process_monitoring_entries",
            ProcessMonitoringEntry,
        ),
        "treated_water_entries": (
            "treated_water_entries",
            TreatedWaterEntry,
        ),
    }

    def parse(self, raw_ocr_document: RawOcrDocument) -> ETPLogbookOCR:
        parsed_document = ETPLogbookOCR(
            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)
            getattr(parsed_document, list_field_name).append(parsed_entry)

        return parsed_document
