"""Post-prediction constraint enforcement.

Mirrors the logical consistency rules from PageStructure.validate_consistency()
in agentic/label_pages.py (lines 117-130). Instead of raising errors, silently
fixes predictions to be logically consistent.
"""

from __future__ import annotations

from typing import Dict

TABLE_SUB_FIELDS = [
    "has_column_headers",
    "table_has_h_ruling_lines",
    "table_has_h_zebra_stripes",
    "table_has_h_whitespace",
    "table_has_v_ruling_lines",
    "table_has_v_whitespace",
    "has_multiple_tables",
]


def enforce_constraints(preds: Dict[str, bool]) -> Dict[str, bool]:
    """Apply logical constraints to raw predictions, returning a corrected copy.

    Rules:
    1. has_hand_filled_form -> has_handwriting must be True
    2. has_empty_form + has_hand_filled_form -> clear has_empty_form
    3. has_multiple_tables -> has_table must be True
    4. has_column_headers -> has_table must be True
    5. not has_table -> all table sub-fields must be False
    """
    out = dict(preds)

    # Rule 1: hand_filled_form implies handwriting
    if out.get("has_hand_filled_form") and not out.get("has_handwriting"):
        out["has_handwriting"] = True

    # Rule 2: empty_form and hand_filled_form are mutually exclusive
    if out.get("has_empty_form") and out.get("has_hand_filled_form"):
        out["has_empty_form"] = False

    # Rule 3: multiple_tables implies has_table
    if out.get("has_multiple_tables") and not out.get("has_table"):
        out["has_table"] = True

    # Rule 4: column_headers implies has_table
    if out.get("has_column_headers") and not out.get("has_table"):
        out["has_table"] = True

    # Rule 5: no table -> clear all table sub-features
    if not out.get("has_table"):
        for field in TABLE_SUB_FIELDS:
            out[field] = False

    return out
