from dataclasses import dataclass
from typing import Literal

from ...dataset import DatasetCfg
from .decoder import Decoder

try:
    from .decoder_splatting_cuda import DecoderSplattingCUDA, DecoderSplattingCUDACfg
except ImportError as exc:
    DecoderSplattingCUDA = None
    CUDA_SPLATTING_IMPORT_ERROR = exc

    @dataclass
    class DecoderSplattingCUDACfg:
        name: Literal["splatting_cuda"]
else:
    CUDA_SPLATTING_IMPORT_ERROR = None

from .decoder_splatting_gsplat import (
    DecoderSplattingGsplat,
    DecoderSplattingGsplatCfg,
)

DECODERS = {
    "splatting_gsplat": DecoderSplattingGsplat,
    "gsplat": DecoderSplattingGsplat,
}

DecoderCfg = DecoderSplattingGsplatCfg | DecoderSplattingCUDACfg

if DecoderSplattingCUDA is not None:
    DECODERS["splatting_cuda"] = DecoderSplattingCUDA


def get_decoder(decoder_cfg: DecoderCfg, dataset_cfg: DatasetCfg) -> Decoder:
    decoder_cls = DECODERS.get(decoder_cfg.name)
    if decoder_cls is None:
        raise ImportError(
            "The 'splatting_cuda' decoder was requested, but its CUDA rasterizer "
            "dependencies could not be imported. Use model/decoder=splatting_gsplat "
            "or install the original CUDA splatting dependencies."
        ) from CUDA_SPLATTING_IMPORT_ERROR
    return decoder_cls(decoder_cfg, dataset_cfg)