from __future__ import annotations

import math
import os
import subprocess

import gradio as gr
import pygltflib
import trimesh


def convert_formats(path_input, target_ext):
    """
    Converts an input 3D model under input path to the target extensions format and returns a path to that file.
    :param path_input: path to user input
    :param target_ext: target extension
    :return: path to the input 3D model stored in target format.
    """
    path_input_base, ext = os.path.splitext(path_input)
    if ext == "." + target_ext:
        return path_input
    path_output = path_input_base + "." + target_ext
    if not os.path.exists(path_output):
        trimesh.load_mesh(path_input).export(path_output)
    return path_output


def add_lights(path_input, path_output):
    """
    Adds directional lights in the horizontal plane to the glb file.
    :param path_input: path to input glb
    :param path_output: path to output glb
    :return: None
    """
    glb = pygltflib.GLTF2().load(path_input)

    N = 3  # default max num lights in Babylon.js is 4
    angle_step = 2 * math.pi / N

    lights_extension = {
        "lights": [
            {"type": "directional", "color": [1.0, 1.0, 1.0], "intensity": 2.0}
            for _ in range(N)
        ]
    }

    if "KHR_lights_punctual" not in glb.extensionsUsed:
        glb.extensionsUsed.append("KHR_lights_punctual")
    glb.extensions["KHR_lights_punctual"] = lights_extension

    light_nodes = []
    for i in range(N):
        angle = i * angle_step
        rotation = [0.0, math.sin(angle / 2), 0.0, math.cos(angle / 2)]
        node = {
            "rotation": rotation,
            "extensions": {"KHR_lights_punctual": {"light": i}},
        }
        light_nodes.append(node)

    light_node_indices = list(range(len(glb.nodes), len(glb.nodes) + N))
    glb.nodes.extend(light_nodes)

    root_node_index = glb.scenes[glb.scene].nodes[0]
    root_node = glb.nodes[root_node_index]
    if hasattr(root_node, "children"):
        root_node.children.extend(light_node_indices)
    else:
        root_node.children = light_node_indices

    glb.save(path_output)


def breathe_new_life_into_3d_model(path_input, prompt, path_output=None):
    """
    Paints the input 3D model using the following method:
    @inproceedings{wang2023breathing,
        title={Breathing New Life into 3D Assets with Generative Repainting},
        author={Wang, Tianfu and Kanakis, Menelaos and Schindler, Konrad and Van Gool, Luc and Obukhov, Anton},
        booktitle={Proceedings of the British Machine Vision Conference (BMVC)},
        year={2023},
        publisher={BMVA Press}
    }
    :param path_input: path to input 3D model
    :param prompt: text description of the expected output
    :param path_output: path to precomputed (cached) output if available
    :return: path to the painted 3D model in glb format with lights for Gradio web component
    """
    if path_output is not None:
        path_output = path_output.name
        path_output_glb_vis = path_output[:-4] + "_vis.glb"
        add_lights(path_output, path_output_glb_vis)
        return path_output_glb_vis

    path_output_dir = path_input + ".output"
    os.makedirs(path_output_dir, exist_ok=True)

    path_input_ply = convert_formats(path_input, "ply")

    cmd = [
        "bash",
        "/repainting_3d_assets/code/scripts/conda_run.sh",
        "/repainting_3d_assets",
        path_input_ply,
        path_output_dir,
        prompt,
    ]
    result = subprocess.run(cmd, env=os.environ, text=True)

    if result.returncode != 0:
        print(f"Output: {result.stdout}")
        print(f"Stderr: {result.stderr}")
        raise RuntimeError("Processing failed")

    path_output_glb = os.path.join(path_output_dir, "model_draco.glb")
    path_output_glb_vis = path_output_glb[:-4] + "_vis.glb"
    add_lights(path_output_glb, path_output_glb_vis)

    return path_output_glb_vis


def run():
    """
    Gradio entry point
    """
    desc = """
        <p align="center">
        <a title="Website" href="https://www.obukhov.ai/repainting_3d_assets" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
            <img src="https://www.obukhov.ai/img/badges/badge-website.svg">
        </a>
        <a title="arXiv" href="https://arxiv.org/abs/2309.08523" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
            <img src="https://www.obukhov.ai/img/badges/badge-pdf.svg">
        </a>
        <a title="Github" href="https://github.com/kongdai123/repainting_3d_assets" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
            <img src="https://img.shields.io/github/stars/kongdai123/repainting_3d_assets?label=GitHub%20%E2%98%85&logo=github&color=C8C" alt="badge-github-stars">
        </a>
        <a title="Social" href="https://twitter.com/antonobukhov1" target="_blank" rel="noopener noreferrer" style="display: inline-block;">
            <img src="https://www.obukhov.ai/img/badges/badge-social.svg" alt="social">
        </a>
        </p>
        <p align="justify">
        Repaint your 3D models with a text prompt, guided by a method from our BMVC'2023 Oral paper 'Breathing New Life 
        into 3D Assets with Generative Repainting'. Simply drop a model into the left pane, specify your repainting 
        preferences, and wait for the outcome (~20 min). Explore precomputed examples at the bottom, or follow the 
        Project Website badge for additional precomputed models and comparison with other repainting techniques.    
        </p>
    """
    demo = gr.Interface(
        title="Repainting 3D Assets",
        description=desc,
        thumbnail="thumbnail.jpg",
        fn=breathe_new_life_into_3d_model,
        inputs=[
            gr.Model3D(
                camera_position=(30.0, 90.0, 3.0),
                elem_classes="viewport",
                label="Input Model",
            ),
            gr.Textbox(
                label="Text Prompt",
            ),
            gr.File(visible=False, label="Cached Output"),
        ],
        outputs=[
            gr.Model3D(
                camera_position=(30.0, 90.0, 3.0),
                elem_classes="viewport",
                label="Repainted Model",
            ),
        ],
        examples=[
            [
                os.path.join(os.path.dirname(__file__), "files/horse.glb"),
                "pastel superhero unicorn",
                os.path.join(
                    os.path.dirname(__file__), "files/horse.glb.output/output.glb"
                ),
            ],
        ],
        cache_examples=True,
        css="""
            .viewport {
                aspect-ratio: 16/9;
            }
        """,
        allow_flagging="never",
    )

    demo.queue().launch(server_name="0.0.0.0", server_port=7860)


if __name__ == "__main__":
    run()