from dataclasses import dataclass, fields
from typing import Any, ClassVar, Mapping, Optional

from schemas_hf.base import DatasetSchema
from schemas_hf.registry import register_schema


@register_schema
@dataclass
class UwunionCelebaAttrsRecord(DatasetSchema):
    DATASET_NAME: ClassVar[str] = "uwunion/celeba-attrs"
    CONFIG_NAME: ClassVar[Optional[str]] = None
    rid: int
    image_id: str
    image_path: str
    # Original HF column name is '5_o_Clock_Shadow'; we expose it as a valid
    # Python identifier and override from_example to remap the key.
    five_o_clock_shadow: bool
    Arched_Eyebrows: bool
    Attractive: bool
    Bags_Under_Eyes: bool
    Bald: bool
    Bangs: bool
    Big_Lips: bool
    Big_Nose: bool
    Black_Hair: bool
    Blond_Hair: bool
    Blurry: bool
    Brown_Hair: bool
    Bushy_Eyebrows: bool
    Chubby: bool
    Double_Chin: bool
    Eyeglasses: bool
    Goatee: bool
    Gray_Hair: bool
    Heavy_Makeup: bool
    High_Cheekbones: bool
    Male: bool
    Mouth_Slightly_Open: bool
    Mustache: bool
    Narrow_Eyes: bool
    No_Beard: bool
    Oval_Face: bool
    Pale_Skin: bool
    Pointy_Nose: bool
    Receding_Hairline: bool
    Rosy_Cheeks: bool
    Sideburns: bool
    Smiling: bool
    Straight_Hair: bool
    Wavy_Hair: bool
    Wearing_Earrings: bool
    Wearing_Hat: bool
    Wearing_Lipstick: bool
    Wearing_Necklace: bool
    Wearing_Necktie: bool
    Young: bool
    image: Any

    @classmethod
    def from_example(cls, example: Mapping[str, Any]) -> "UwunionCelebaAttrsRecord":  # type: ignore[override]
        # Map the HF column '5_o_Clock_Shadow' into the dataclass field
        # 'five_o_clock_shadow' before delegating to the base logic.
        mapping: dict[str, Any] = dict(example)
        if "5_o_Clock_Shadow" in mapping:
            mapping["five_o_clock_shadow"] = mapping["5_o_Clock_Shadow"]
        # Ensure all init fields are present.
        field_names = [f.name for f in fields(cls) if f.init]
        missing = [name for name in field_names if name not in mapping]
        if missing:
            raise KeyError(f"Missing fields for {cls.__name__}: {', '.join(missing)}")
        values = {name: mapping[name] for name in field_names}
        return cls(**values)
