from __future__ import annotations
from typing import TYPE_CHECKING, ClassVar
from dimetra.geo.length_unit import length_unit_from_token
from h5kit import codec
from h5kit.codec import H5CodecLike, H5CodecMixin
from dimetra.geo.length import Length
from dimetra.h5.codec.length_codec_layout import LENGTH_CODEC_LAYOUT

if TYPE_CHECKING:
    from h5kit.io import H5Reader, H5Writer


@codec
class LengthCodec(H5CodecMixin, H5CodecLike[Length]):
    PY_TYPE: ClassVar[type[Length]] = Length
    SCHEMA: ClassVar[str] = "allytools.units.Length"
    VER: ClassVar[int] = 1
    LAYOUT: ClassVar = LENGTH_CODEC_LAYOUT.h5_layout
    DEFAULT_UNIT: ClassVar[str] = "mm"

    @classmethod
    def write(cls, w: H5Writer, obj: Length) -> None:
        cls.clear_before_write(w)
        unit_obj = getattr(obj, "_original_unit", None)
        unit_token = (
            getattr(unit_obj, "name", cls.DEFAULT_UNIT)
            if unit_obj is not None
            else cls.DEFAULT_UNIT)
        w.attr(LENGTH_CODEC_LAYOUT.value_mm, float(obj.value_mm))
        w.attr(LENGTH_CODEC_LAYOUT.unit, unit_token)

    @classmethod
    def read(cls, r: H5Reader) -> Length:
        mm = r.require_attr(LENGTH_CODEC_LAYOUT.value_mm)
        token = str(r.attr(LENGTH_CODEC_LAYOUT.unit, cls.DEFAULT_UNIT))
        unit = length_unit_from_token(token)
        from dimetra.geo.length import Length
        value_in_unit = mm / unit.factor
        return Length(value_in_unit, unit)