import copy
from typing import Optional

from PIL import Image

from .single_image_convsation import SingleImageConvDatasetMixin


class SingleImageInteractive(SingleImageConvDatasetMixin):
    _printed_sample = True

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.image: Optional[Image.Image] = None
        self.roles = ('human', 'gpt')
        self.boxes = []
        self.points = []
        self.raw_conv = []
        self.conversations = []

    def set_image(self, image: Image.Image):
        assert self.image is None, f"{image}"
        self.image = image

    def append_message(self, role: str, message: str, *, boxes=None, points=None, boxes_seq=None, points_seq=None):
        """Append a new message."""
        assert role in self.roles

        def convert_idx(objs_seq, objs_value, get_obj_idx_func):
            if objs_seq is None:
                return None
            ret = []
            for objs_idx in objs_seq:
                new_objs_idx = []
                for idx in objs_idx:
                    new_idx = get_obj_idx_func(objs_value[idx])
                    new_objs_idx.append(new_idx)
                ret.append(tuple(new_objs_idx))
            return tuple(ret)

        boxes_seq = convert_idx(boxes_seq, boxes, self._get_box_idx)
        points_seq = convert_idx(points_seq, points, self._get_point_idx)

        if self.image is not None:
            previous_message_has_image_placeholder = any(
                '<image>' in item['value'] for item in self.conversations
            )
            if not previous_message_has_image_placeholder and '<image>' not in message:
                message = '<image> ' + message
            if previous_message_has_image_placeholder and '<image>' in message:
                message = message.replace('<image>', '')

        self.conversations.append(
            {
                'from': role,
                'value': message,
                'boxes_seq': copy.deepcopy(boxes_seq),
                'points_seq': copy.deepcopy(points_seq),
            }
        )

    def get_raw_item(self, index=None):
        ret = copy.deepcopy({
            'image': self.image,
            'target': {
                'boxes': self.boxes,
                'points': self.points,
            },
            'conversations': self.conversations,
        })
        assert ret['conversations'][0]['from'] == self.roles[0]
        if ret['conversations'][-1]['from'] == self.roles[0]:
            ret['conversations'].append(
                {
                    'from': self.roles[1],
                    'value': '',
                }
            )
        return ret

    def to_model_input(self):
        item = self.__getitem__(0)
        ret = {'input_ids': item['input_ids'].unsqueeze(0).cuda()}
        if 'image' in item and item['image'] is not None:
            ret['images'] = item['image'].unsqueeze(0).cuda()
        else:
            ret['images'] = None
        return ret

    def to_gradio_chatbot_new_messages(self):
        conv = self.__getitem__(0, return_conv=True)
        new_messages = conv.messages[-2:]
        ret_messages = []
        for r, m in new_messages:
            nm = m.replace('<im_patch>', '').replace('<im_end>', '').replace('<im_start>', '<image>')
            ret_messages.append((r, nm))
        return ret_messages

    def _get_box_idx(self, box):
        assert isinstance(box, (tuple, list)), f"{type(box)}"
        assert isinstance(box[0], (int, float)), f"{type(box[0])}"
        assert len(box) == 4
        box = tuple(box)
        if box not in self.boxes:
            self.boxes.append(box)
            return len(self.boxes) - 1
        else:
            return self.boxes.index(box)

    def _get_point_idx(self, point):
        assert isinstance(point, (tuple, list))
        assert isinstance(point[0], (int, float))
        assert len(point) == 2
        point = tuple(point)
        if point not in self.points:
            self.points.append(tuple(point))
            return len(self.points) - 1
        else:
            return self.points.index(point)

    def __len__(self):
        return 1