from decimal import Decimal
from typing import Annotated

from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator

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

class AcfEntry(TimedLogEntry):
    appearance_and_odour: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["Nm/Noo"]), ocr_default="Nm/Noo")] = None
    line_number: OcrField[str] | None = None
    m_alkalinity: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=85))] = None
    pH: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=6.5, max=8.5))] = None
    total_chlorine: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=0))] = None
    total_hardness:  Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=250))] = None
    turbidity: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=1))] = None


class UfEntry(TimedLogEntry):
    turbidity: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=1))] = None


class LeadEntry(TimedLogEntry):
    appearance_and_odour: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["Nm/Noo"]), ocr_default="Nm/Noo")] = None
    pH: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=6.5, max=8.5))] = None
    tds: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=250), required=True)] = None
    turbidity: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=1))] = None


class LagEntry(TimedLogEntry):
    appearance_and_odour: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["Nm/Noo"]), ocr_default="Nm/Noo")] = None
    line_number: OcrField[str] | None = None
    pH: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=6.5, max=8.5))] = None
    tds:  Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=250))] = None
    turbidity: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=1))] = None


class RoEntry(TimedLogEntry):
    appearance_odour_taste: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["Nm/Noo/Not"]), ocr_default="Nm/Noo/Not")] = None
    feed_pH: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=6.5, max=8.5))] = None
    feed_tds:  Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=250), required=True)] = None
    permeate_pH: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=6.5, max=8.5))] = None
    permeate_tds:  Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=85), required=True)] = None
    turbidity: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=1))] = None


class SandFilterEntry(TimedLogEntry):
    appearance_and_odour: Annotated[OcrField[str] | None, FieldSpec(ocr_spec=OptionSpec[str](options=["Nm/Noo"]), ocr_default="Nm/Noo")] = None
    free_chlorine:  Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=1))] = None
    line_number: OcrField[str] | None = None
    m_alkalinity: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=85))] = None
    p_alkalinity:  Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=10))] = None
    pH: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=6.5, max=8.5))] = None
    turbidity: Annotated[OcrField[Decimal] | None, FieldSpec(ocr_spec=RangeSpec[Decimal](min=0, max=1))] = None

TABLE_ORDER = [
    "sand_filter_entry",
    "uf_entry",
    "acf_entry",
    "lead_entry",
    "lag_entry",
    "ro_entry",
]


class WTPLeftPageLogRow(BaseModel):
    sand_filter_entry: SandFilterEntry | None = None
    uf_entry: UfEntry | None = None
    lead_entry: LeadEntry | None = None
    lag_entry: LagEntry | None = None
    ro_entry: RoEntry | None = None
    acf_entry: AcfEntry | None = None
    # Set by compute_chronology() after construction -- a row with an out-of-order time
    # between filtration stages is kept (not discarded) but flagged for reviewer attention.
    chronology_flag: bool = False
    chronology_issues: list[str] = Field(default_factory=list)

    @model_validator(mode="after")
    def validate_entry_times_present(self):
        # A row can only be aligned across the 6 filtration stages if every stage has a
        # time to compare -- that's genuinely unrepairable, unlike ordering (see
        # compute_chronology), so it still causes this row to be dropped by the caller.
        for field_name in TABLE_ORDER:
            self._entry_time_minutes(field_name, getattr(self, field_name))
        return self

    def compute_chronology(self) -> None:
        """Flags (rather than discards) a row where a later filtration stage's time isn't
        strictly after -- and within 15 minutes of -- the previous stage's time."""
        previous_time_minutes: int | None = None
        previous_field_name: str | None = None
        issues: list[str] = []

        for field_name in TABLE_ORDER:
            entry_time = self._entry_time_minutes(field_name, getattr(self, field_name))

            if previous_time_minutes is not None:
                if entry_time <= previous_time_minutes:
                    issues.append(
                        f"{field_name}.time must be after {previous_field_name}.time"
                    )
                elif entry_time - previous_time_minutes >= 15:
                    issues.append(
                        f"{field_name}.time must be less than 15 minutes after "
                        f"{previous_field_name}.time"
                    )

            previous_time_minutes = entry_time
            previous_field_name = field_name

        self.chronology_issues = issues
        self.chronology_flag = bool(issues)

    @staticmethod
    def _entry_time_minutes(field_name: str, entry: TimedLogEntry | None) -> int:
        if entry is None or entry.time is None or entry.time.value is None:
            raise ValueError(f"{field_name}.time is required")

        return entry.time.value.hour * 60 + entry.time.value.minute


class WTPLeftPageOCR(LogbookOCR):
    log_rows: list[WTPLeftPageLogRow] = Field(default_factory=list)
    sand_filter_entries: list[SandFilterEntry] = Field(default_factory=list)
    uf_entries: list[UfEntry] = Field(default_factory=list)
    lead_entries: list[LeadEntry] = Field(default_factory=list)
    lag_entries: list[LagEntry] = Field(default_factory=list)
    ro_entries: list[RoEntry] = Field(default_factory=list)
    acf_entries: list[AcfEntry] = Field(default_factory=list)
    model_config = ConfigDict(populate_by_name=True)


class WTPLeftPageParser(OcrParser):
    _ENTRY_TYPES: dict[str, tuple[str, type[TimedLogEntry]]] = {
        "ACF": ("acf_entries", AcfEntry),
        "UF": ("uf_entries", UfEntry),
        "lead": ("lead_entries", LeadEntry),
        "lag": ("lag_entries", LagEntry),
        "ro": ("ro_entries", RoEntry),
        "sand_filter_log_entry": ("sand_filter_entries", SandFilterEntry),
    }

    def parse(self, raw_ocr_document: RawOcrDocument) -> WTPLeftPageOCR:

        parsed_document = WTPLeftPageOCR(
            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)

        parsed_document.log_rows = self._build_log_rows(parsed_document)
        return parsed_document

    def _build_log_rows(
        self, parsed_document: WTPLeftPageOCR
    ) -> list[WTPLeftPageLogRow]:
        entry_lists = {
            log_row_field_name: getattr(parsed_document, list_field_name)
            for log_row_field_name, list_field_name in (
                self._log_row_list_fields().items()
            )
        }
        # The 6 stage tables are extracted independently, so their entry counts routinely
        # differ by a row or two (one table's last row gets missed, an extra spurious row
        # gets picked up, etc.) -- that used to zero out log_rows for the *entire* page,
        # discarding every stage's data (including e.g. TDS, which only appears in
        # lead/ro entries) even when most rows lined up fine. Aligning on the shortest
        # list still assumes index N means "the same physical row" across all 6 stages,
        # which only holds for the leading rows before whichever stage diverges first --
        # but that recovers real data instead of dropping the whole page over a routine
        # single-row extraction miss.
        entry_count = min(len(entries) for entries in entry_lists.values())
        log_rows: list[WTPLeftPageLogRow] = []
        for row_index in range(entry_count):
            try:
                row = WTPLeftPageLogRow(
                    **{
                        field_name: entries[row_index]
                        for field_name, entries in entry_lists.items()
                    }
                )
            except ValidationError:
                # A stage genuinely missing its time can't be aligned or compared at all --
                # drop just this row, not the whole page (unlike an ordering violation,
                # which compute_chronology flags below instead of raising).
                continue

            row.compute_chronology()
            log_rows.append(row)

        return log_rows

    @staticmethod
    def _log_row_list_fields() -> dict[str, str]:
        return {
            "sand_filter_entry": "sand_filter_entries",
            "uf_entry": "uf_entries",
            "acf_entry": "acf_entries",
            "lead_entry": "lead_entries",
            "lag_entry": "lag_entries",
            "ro_entry": "ro_entries",
        }