import dataclasses

from flask import jsonify, request

import gazebo_ctrl
import scene_plugin

from .gen_models import GraspModelGenerator


@dataclasses.dataclass
class BoxModel:
    length: float
    weight: float
    height: float
    tag_id: int
    mass: float
    type: str = 'box'


@dataclasses.dataclass
class TableModel:
    height: float
    type: str = 'table'


@dataclasses.dataclass
class TagModel:
    tag_id: int
    type: str = 'tag'


class Plugin(scene_plugin.ScenePluginBase):
    def __init__(self, scene_name):
        super().__init__(scene_name)
        self.model_dict = {}
        self.spawner = GraspModelGenerator()
        self.gz_ctrl = gazebo_ctrl.GazeboROSController()

    def post_scene_model_spawn(self):
        print(222)
        try:
            # 解析参数
            data = request.get_json(force=True)
            name = data.get('name')
            type_ = data.get('type')
            pose = data.get('pose')
            position = pose.get('position')
            orientation = pose.get('orientation')
            tag_id = data.get('tag_id', 0)
            mass = data.get('mass', 1)
            size = data.get('size')
            if not name or not type_ or pose is None:
                raise ValueError('name, type, and pose are required')
            # 生成模型
            if type_ == 'box':
                if not size or not all(k in size for k in ('l', 'h', 'w')):
                    raise ValueError('size with l, h, w is required for box')
                model = BoxModel(
                    length=size['l'],
                    height=size['h'],
                    weight=size['w'],
                    tag_id=tag_id,
                    mass=mass
                )
                ret, msg = self.spawner.generate_box_with_apriltag(
                    name,
                    l=size['l'],
                    w=size['w'],
                    h=size['h'],
                    mass=mass,
                    tag_id=tag_id,
                    x=position.get('x', 0),
                    y=position.get('y', 0),
                    z=position.get('z', 0),
                    roll=orientation.get('roll', 0),
                    pitch=orientation.get('pitch', 0),
                    yaw=orientation.get('yaw', 0)
                )
                if ret == False:
                    raise RuntimeError(msg)
            elif type_ == 'table':
                model = TableModel(
                    height=pose.get('z', 1)
                )
                self.spawner.generate_table(
                    name,
                    x=position.get('x', 0),
                    y=position.get('y', 0),
                    height=position.get('z', 0),
                    roll=orientation.get('roll', 0),
                    pitch=orientation.get('pitch', 0),
                    yaw=orientation.get('yaw', 0)
                )
            elif type_ == 'tag':
                model = TagModel(
                    tag_id=tag_id
                )
                self.spawner.generate_apriltag(
                    name,
                    tag_id=tag_id,
                    x=position.get('x', 0),
                    y=position.get('y', 0),
                    z=position.get('z', 0),
                    roll=orientation.get('roll', 0),
                    pitch=orientation.get('pitch', 0),
                    yaw=orientation.get('yaw', 0)
                )
            else:
                raise ValueError('Invalid type')
            # 注册模型
            self.model_dict[name] = model
        except Exception as e:
            print(str(e))
            return str(e), 400
        return 'OK', 200

    def get_scene_model_state(self):
        try:
            res = self.gz_ctrl.get_model_state_f(self.model_dict.keys())
            for item in res:
                model = self.model_dict.get(item['name'])
                if model.type == 'box':
                    item['tag_id'] = model.tag_id
                    item['mass'] = model.mass
                    item['size'] = {'l': model.length, 'w': model.weight, 'h': model.height}
                elif model.type == 'table':
                    item['pose']['position']['z'] = model.height
                elif model.type == 'tag':
                    item['tag_id'] = model.tag_id
            return jsonify(res), 200
        except Exception as e:
            return str(e), 400
    
    def post_scene_model_state(self):
        try:
            # 解析参数
            data = request.get_json(force=True)
            name = data.get('name')
            type_ = data.get('type')
            pose = data.get('pose')
            position = pose.get('position')
            orientation = pose.get('orientation')
            tag_id = data.get('tag_id', 0)
            mass = data.get('mass', 0)
            size = data.get('size')
            if not name or not type_ or pose is None:
                raise ValueError('name, type, and pose are required')
            # 修改模型
            model = self.model_dict[name]
            if model.type != type_:
                raise ValueError(f'Cannot change model type from {self.model_dict[name].type} to {type_}')
            # 修改pose
            pass
            # 修改独特属性
            if type_ == 'box':
                if mass:
                    model.mass = mass
                if tag_id:
                    model.tag_id = tag_id
            elif type_ == 'table':
                height = position.get('z', 0)
                if height:
                    model.height = height
            elif type_ == 'tag':
                if tag_id:
                    model.tag_id = tag_id
            else:
                raise ValueError('Invalid type')
        except Exception as e:
            return str(e), 400
        return 'OK', 200

if __name__ == '__main__':
    pass