import os
import random
import carb
import omni.usd
import omni.kit.commands
from pxr import Usd, Sdf, UsdGeom, Gf, Tf, UsdShade, UsdLux


def create_test_stage():
    stage = omni.usd.get_context().get_stage()
    rootname = stage.GetDefaultPrim().GetPath().pathString
    prim_list1 = []
    prim_list2 = []

    # create Looks folder
    omni.kit.commands.execute(
        "CreatePrim", prim_path="{}/Looks".format(rootname), prim_type="Scope", select_new_prim=False
    )

    # create material 1
    mtl_created_list = []
    omni.kit.commands.execute(
        "CreateAndBindMdlMaterialFromLibrary",
        mdl_name="OmniPBR.mdl",
        mtl_name="OmniPBR",
        mtl_created_list=mtl_created_list,
    )
    mtl_path1 = mtl_created_list[0]

    # create material 2
    mtl_created_list = []
    omni.kit.commands.execute(
        "CreateAndBindMdlMaterialFromLibrary",
        mdl_name="OmniGlass.mdl",
        mtl_name="OmniGlass",
        mtl_created_list=mtl_created_list,
    )
    mtl_path2 = mtl_created_list[0]

    # create prims & bind material
    random_list = {}
    samples = 100
    for index in random.sample(range(samples), int(samples / 2)):
        random_list[index] = 0

    strengths = [UsdShade.Tokens.strongerThanDescendants, UsdShade.Tokens.weakerThanDescendants]
    for index in range(samples):
        prim_path = omni.usd.get_stage_next_free_path(stage, "{}/TestCube_{}".format(rootname, index), False)
        omni.kit.commands.execute("CreatePrim", prim_path=prim_path, prim_type="Cube")
        if index in random_list:
            omni.kit.commands.execute(
                "BindMaterial", prim_path=prim_path, material_path=mtl_path1, strength=strengths[index & 1]
            )
            prim_list1.append(prim_path)
        elif (index & 3) != 0:
            omni.kit.commands.execute(
                "BindMaterial", prim_path=prim_path, material_path=mtl_path2, strength=strengths[index & 1]
            )
            prim_list2.append(prim_path)

    return [(mtl_path1, prim_list1), (mtl_path2, prim_list2)]


def create_and_select_cube():
    stage = omni.usd.get_context().get_stage()
    rootname = stage.GetDefaultPrim().GetPath().pathString
    prim_path = "{}/Cube".format(rootname)

    # create Looks folder
    omni.kit.commands.execute(
        "CreatePrim", prim_path=prim_path,
        prim_type="Cube", select_new_prim=False,
        attributes={UsdGeom.Tokens.size: 100}
    )

    return prim_path
