import json
from agent import *
from argparse import ArgumentParser
from get_image_from_glb import *
import os
import base64
import pprint
import time
import random


class subPart(BaseModel):
    name: str
    color: str
    shape: str
    size: str
    material: str
    functionality: str
    texture: str


class ObjDescFormat(BaseModel):
    raw_description: str = Field(description="the name of the object,without index and '_'")
    wholePart: subPart = Field(description="the object as a whole")
    subParts: List[subPart] = Field(
        description="the deformable subparts of the object.If the object is not deformable, leave empty here")
    description: List[str] = Field(description="several different text descriptions describing this same object here")
    # val_description:List[str]=Field(description="similar to descriptions, used for validation")


with open("./_generate_object_prompt.txt", "r") as f:
    system_prompt = f.read()


def save_json(save_dir, glb_file_name, ObjDescResult):
    os.makedirs(save_dir, exist_ok=True)
    # Remove .glb extension from the filename
    base_name = glb_file_name.replace(".glb", "")
    save_path = f"{save_dir}/{base_name}.json"

    # Get all descriptions
    all_descriptions = ObjDescResult.description.copy()
    all_descriptions.sort(key=len)
    # Randomly select 5 indices for validation set
    val_indices = random.sample(range(len(all_descriptions)), 3)

    # Separate validation and training descriptions based on indices
    shuffle_val = [all_descriptions[i] for i in val_indices]
    shuffle_train = [all_descriptions[i] for i in range(len(all_descriptions)) if i not in val_indices]

    # Sort both validation and training descriptions by character length
    shuffle_val.sort(key=len)
    shuffle_train.sort(key=len)

    # 将字典保存为 JSON 文件
    desc_dict = {
        "raw_description": ObjDescResult.raw_description,
        "seen": shuffle_train,
        "unseen": shuffle_val,
    }
    with open(save_path, "w", encoding="utf-8") as file:
        json.dump(desc_dict, file, ensure_ascii=False, indent=4)
        print(json.dumps(desc_dict, indent=2, ensure_ascii=False))


def save_image(save_dir, glb_file_name, imgstr):
    os.makedirs(save_dir, exist_ok=True)
    save_image_path = f"{save_dir}/{glb_file_name}.png"
    with open(save_image_path, "wb") as f:
        # Convert the Base64 string to bytes before writing
        img_data = base64.b64decode(imgstr)
        f.write(img_data)


def make_prompt_generate(imgStr, object_name):
    messages = [
        {
            "role": "system",
            "content": system_prompt
        },
        {
            "role":
            "user",
            "content": [
                {
                    "type": "text",
                    "text": f"THE OBJECT IS A {object_name}"
                },
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/png;base64,{imgStr}"
                    },
                },
            ],
        },
    ]
    result = generate(messages, ObjDescFormat)
    result_dict = result.model_dump()
    print(
        json.dumps(
            {
                "wholePart": result_dict["wholePart"],
                "subParts": result_dict["subParts"],
            },
            indent=2,
            ensure_ascii=False,
        ))
    return result


def generate_obj_description(object_name, glb_file_name):
    time_start = time.time()
    object_file_path = f"../assets/objects/{object_name}/visual/{glb_file_name}"
    save_dir = f"./objects_description/{object_name}"
    result_img_path = f"{save_dir}/{glb_file_name}.png"
    if not os.path.exists(result_img_path):
        imgstr = get_image_from_glb(object_file_path)
        print(f"{object_name} {glb_file_name} saving image", time.time() - time_start)
        time_start = time.time()
        save_image(save_dir, glb_file_name, imgstr)
    else:
        print(
            f'{object_name} {glb_file_name} using existing image: {result_img_path}. If errors like "Message: Invalid image data." occurs, please delete the image and rerun the script'
        )
        with open(result_img_path, "rb") as f:
            imgstr = base64.b64encode(f.read()).decode("utf-8")
    print(f"{object_name} {glb_file_name} start generating", time.time() - time_start)
    time_start = time.time()
    result = make_prompt_generate(imgstr, object_name)
    print(
        f"{object_name} {glb_file_name} generated {len(str(result.model_dump()))} descriptions ",
        time.time() - time_start,
    )
    save_json(save_dir, glb_file_name, result)


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("object_name", type=str, nargs="?", default=None, help="Object name to process")
    parser.add_argument("--index", type=int, default=None, help="Specific object index to process")
    parser.add_argument("--store_png", action="store_true", help="Store PNG files after generation")
    usr_args = parser.parse_args()

    object_name = usr_args.object_name
    object_index = usr_args.index
    clear_png = not usr_args.store_png

    if object_name is None:  # process all objects
        objects_dir = "../assets/objects"
        results_dir = "./objects_description"
        for object_name in sorted(os.listdir(objects_dir)):
            parts = object_name.split("_")
            if not (len(parts) == 2):
                continue
            object_dir = os.path.join(objects_dir, object_name)
            if os.path.isdir(object_dir):
                visual_dir = os.path.join(object_dir, "visual")
                if os.path.exists(visual_dir):
                    print(f"Processing object: {object_name}")
                    glb_files = [file for file in os.listdir(visual_dir) if file.endswith(".glb")]
                    for glb_file in sorted(glb_files):
                        if os.path.exists(os.path.join(
                                results_dir,
                                object_name,
                                glb_file.replace(".glb", ".json"),
                        )):
                            continue
                        generate_obj_description(object_name, glb_file)
                        if clear_png:
                            png_path = (f"./objects_description/{object_name}/{glb_file}.png")
                            if os.path.exists(png_path):
                                os.remove(png_path)
                                print(f"Deleted: {png_path}")
    elif object_index is None:  # all type for specific object
        folder_path = f"../assets/objects/{object_name}/visual"
        files_and_folders = os.listdir(folder_path)
        glb_files = [file for file in files_and_folders if file.endswith(".glb")]
        for glb_file in glb_files:
            generate_obj_description(object_name, glb_file)
            if clear_png:
                png_path = f"./objects_description/{object_name}/{glb_file}.png"
                if os.path.exists(png_path):
                    os.remove(png_path)
                    print(f"Deleted: {png_path}")
    else:  # specific object and index
        generate_obj_description(object_name, f"base{object_index}.glb")
        if clear_png:
            png_path = f"./objects_description/{object_name}/base{object_index}.glb.png"
            if os.path.exists(png_path):
                os.remove(png_path)
                print(f"Deleted: {png_path}")
