from ..smp import *
from .dataset_config import img_root_map
from abc import abstractmethod


class CustomPrompt:

    @abstractmethod
    def use_custom_prompt(self, dataset):
        raise NotImplementedError

    @abstractmethod
    def build_prompt(self, line, dataset):
        raise NotImplementedError

    def dump_image(self, line, dataset):
        ROOT = LMUDataRoot()
        assert isinstance(dataset, str)
        img_root = osp.join(ROOT, 'images', img_root_map[dataset] if dataset in img_root_map else dataset)
        os.makedirs(img_root, exist_ok=True)
        if isinstance(line['image'], list):
            tgt_path = []
            assert 'image_path' in line
            for img, im_name in zip(line['image'], line['image_path']):
                path = osp.join(img_root, im_name)
                if not read_ok(path):
                    decode_base64_to_image_file(img, path)
                tgt_path.append(path)
        else:
            tgt_path = osp.join(img_root, f"{line['index']}.jpg")
            if not read_ok(tgt_path):
                decode_base64_to_image_file(line['image'], tgt_path)
        return tgt_path
