from __future__ import annotations

from dataclasses import dataclass


_MATRIX = {1: "bt709", 5: "smpte170m", 6: "smpte170m", 9: "bt2020nc", 10: "bt2020nc"}
_PRIMARIES = {1: "bt709", 9: "bt2020", 12: "smpte432"}
_TRANSFER = {1: "bt709", 16: "smpte2084", 18: "arib-std-b67"}
_RANGE = {0: "tv", 1: "pc"}


@dataclass(frozen=True)
class ColorMetadata:
    matrix: int | None
    primaries: int | None
    transfer: int | None
    color_range: int | None

    @property
    def colorspace(self) -> str | None:
        return _MATRIX.get(self.matrix)

    def ffmpeg_args(self) -> list[str]:
        values = [
            ("-colorspace", self.colorspace),
            ("-color_primaries", _PRIMARIES.get(self.primaries)),
            ("-color_trc", _TRANSFER.get(self.transfer)),
            ("-color_range", _RANGE.get(self.color_range)),
        ]
        return [part for option, value in values if value is not None for part in (option, value)]

    def x264_args(self) -> list[str]:
        values = [
            ("--colormatrix", self.colorspace),
            ("--colorprim", _PRIMARIES.get(self.primaries)),
            ("--transfer", _TRANSFER.get(self.transfer)),
            ("--range", _RANGE.get(self.color_range)),
        ]
        return [part for option, value in values if value is not None for part in (option, value)]

    def x265_params(self) -> str:
        values = [
            ("colorprim", _PRIMARIES.get(self.primaries)),
            ("transfer", _TRANSFER.get(self.transfer)),
            ("colormatrix", self.colorspace),
            ("range", "full" if self.color_range == 1 else "limited" if self.color_range == 0 else None),
        ]
        return ":".join(f"{name}={value}" for name, value in values if value is not None)

    def mkvmerge_args(self) -> list[str]:
        values = [
            ("--color-matrix-coefficients", self.matrix),
            ("--color-primaries", self.primaries),
            ("--color-transfer-characteristics", self.transfer),
            ("--color-range", {0: 1, 1: 2}.get(self.color_range)),
        ]
        return [part for option, value in values if value is not None for part in (option, f"0:{value}")]

    def hevc_bitstream_filter(self) -> str | None:
        values = [
            ("colour_primaries", self.primaries),
            ("transfer_characteristics", self.transfer),
            ("matrix_coefficients", self.matrix),
            ("video_full_range_flag", self.color_range),
        ]
        params = ":".join(f"{name}={value}" for name, value in values if value is not None)
        return f"hevc_metadata={params}" if params else None

    def expected_probe_values(self) -> dict[str, str | None]:
        return {
            "matrix": {1: "Rec709", 5: "Rec601", 6: "Rec601", 9: "Rec2020", 10: "Rec2020"}.get(self.matrix),
            "primaries": _PRIMARIES.get(self.primaries),
            "transfer": _TRANSFER.get(self.transfer),
            "range": _RANGE.get(self.color_range),
        }