from __future__ import annotations
from typing import ClassVar
from dimetra.geo.angle import Angle
from dimetra.geo.angle_unit import angle_unit_from_token
from h5kit import codec
from h5kit.codec import H5CodecLike, H5CodecMixin
from h5kit.io import H5Reader, H5Writer
from dimetra.h5.codec.angle_codec_layout import ANGLE_CODEC_LAYOUT


@codec
class AngleCodec(H5CodecMixin, H5CodecLike[Angle]):
    PY_TYPE: ClassVar[type[Angle]] = Angle
    SCHEMA: ClassVar[str] = "allytools.units.Angle"
    VER: ClassVar[int] = 1
    LAYOUT: ClassVar = ANGLE_CODEC_LAYOUT.h5_layout
    DEFAULT_UNIT: ClassVar[str] = "rad"

    @classmethod
    def write(cls, w: H5Writer, obj: Angle) -> 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(ANGLE_CODEC_LAYOUT.value_rad, float(obj.value_rad))
        w.attr(ANGLE_CODEC_LAYOUT.unit, unit_token)

    @classmethod
    def read(cls, r: H5Reader) -> Angle:
        rad = r.require_attr(ANGLE_CODEC_LAYOUT.value_rad)
        token = str(r.attr(ANGLE_CODEC_LAYOUT.unit, cls.DEFAULT_UNIT))
        unit = angle_unit_from_token(token)
        value_in_unit = rad / unit.factor
        return Angle(value_in_unit, unit)