"""segments → tensor_dict 的多模态转换：每段路由到对应模态 processor，产物合并。"""

from __future__ import annotations

import torch

from deepinfer.ve.processor.common import BaseProcessor, ImageInput, SegmentInput, TextLengthInput, merge_tensor_dicts


class MMProcessor:
    def __init__(
        self,
        text_token_processor: BaseProcessor = None,
        image_processor: BaseProcessor = None,
    ) -> None:
        self.text_token_processor = text_token_processor
        self.image_processor = image_processor
        self.processors_map = {
            TextLengthInput: text_token_processor,
            ImageInput: image_processor,
        }

    def convert_segments_to_tensor_dict(self, seg: SegmentInput) -> dict:
        """每个分段调用该模态对应的 processor，同名 tensor 合并。"""
        # 空段种子保证零段 / 全缓存 prompt 也有稳定的键集合（mm_token_type 恒产出）
        tensor_dicts: list[dict] = [{"mm_token_type": torch.zeros(0, dtype=torch.int32)}]
        for segment in seg.segments:
            processor: BaseProcessor = self.processors_map[type(segment)]
            tensor_dicts.append(processor.process(segment))
        return merge_tensor_dicts(tensor_dicts)