import bpy
import gradio as gr
import math
import os
import random
from diffusers.utils import load_image
from PIL import Image
import PIL.Image
import tqdm

os.putenv("XDG_RUNTIME_DIR", '.')
os.environ['XDG_RUNTIME_DIR'] = '.'
fin_rend = './rendered_scene.png'
fin_blend = './exp_scene.blend'
fin_glb = './exp_scene.glb'
fin_anim = './animation.mp4'
img = bpy.data.images.load('./STScI-01G8H1NK4W8CJYHF2DDFD1W0DQ.png')
earth_texture = bpy.data.images.load('./ear0xuu2.jpg')
mars_texture = bpy.data.images.load('mar0kuu2.jpg')
apol=[]
print(bpy.app.version_string)
def add_planet(name, radius, mass, distance_from_sun):
    material = bpy.data.materials.new(name="Mymarsmat"+str(random.randint(1,36)))
    material.use_nodes = True
    image_texture_node = material.node_tree.nodes.new('ShaderNodeTexImage')
    if name=="mars":
        image_texture_node.image = bpy.data.images.load("mar0kuu2.jpg")
    else:
        image_texture_node.image = bpy.data.images.load("ear0xuu2.jpg")
    bsdf_node = material.node_tree.nodes.get('Principled BSDF')
    material.node_tree.links.new(bsdf_node.inputs['Base Color'], image_texture_node.outputs['Color'])
    bpy.ops.mesh.primitive_uv_sphere_add(radius=radius, location=(distance_from_sun * random.uniform(-8, 8), random.uniform(-8, 8), random.uniform(-8, 3)))
    planet = bpy.context.active_object
    planet.name = name
    planet.data.materials.append(material)
    planet["mass"] = mass
    bpy.context.view_layer.objects.active = planet
    planet.select_set(True)
    planet.animation_data_create()
    planet.animation_data.action = bpy.data.actions.new(name="RotationAction_"+planet.name)
    start_frame = 1
    duration=2
    end_frame = int(duration * 10)
    planet.rotation_mode = 'XYZ'
    planet.keyframe_insert(data_path="rotation_euler", index=2, frame=start_frame)
    
    for frame in range(start_frame, end_frame + 1):
        ##angle = (frame / end_frame) * 2 * math.pi
        ##x = radius * math.cos(angle)
        ##y = planet.location.y
        ##z = radius * math.sin(angle)
        ##planet.location = (x, y, z)
        ##planet.keyframe_insert(data_path="location", index=-1, frame=frame)
        angle = math.radians(frame * 18)
        planet.rotation_euler = (0, 0, angle)
        planet.keyframe_insert(data_path="rotation_euler", index=-1, frame=frame)
    ##return planet

def create_planet_flyby_animation(planet, radius, duration):
    bpy.context.view_layer.objects.active = planet
    planet.select_set(True)
    planet.animation_data_create()
    planet.animation_data.action = bpy.data.actions.new(name="OrbitAction")
    start_frame = 1
    end_frame = int(duration * 10)
    for frame in range(start_frame, end_frame + 1):
        angle = (frame / end_frame) * 2 * math.pi
        x = radius * math.cos(angle)
        y = planet.location.y
        z = radius * math.sin(angle)
        planet.location = (x, y, z)
        planet.keyframe_insert(data_path="location", index=-1, frame=frame)

def add_animated_space_rocks(name):
    bpy.context.scene.render.fps = 10
    duration = 2
    radius=2
    
    bpy.ops.mesh.primitive_uv_sphere_add(radius=random.uniform(0.03, 0.1), location=(random.uniform(-10, 10), random.uniform(-10, 10), random.uniform(-4, 9)))
    space_rock = bpy.context.active_object
    space_rock.name = ""+name+"_"+str(random.randint(1,100))+""
    space_rock.keyframe_insert(data_path="location", frame=1)
    space_rock.location = (random.uniform(-10, 10), random.uniform(-10, 10), random.uniform(-10, 10))
    space_rock.keyframe_insert(data_path="location", frame=20)
    bpy.context.view_layer.objects.active = space_rock
    space_rock.select_set(True)
        

def create_particle_effects_animation():
    ##wip
    bpy.context.scene.render.fps = 10
    duration = 2

    bpy.ops.mesh.primitive_uv_sphere_add(radius=random.uniform(0.05, 0.10), location=(random.uniform(-1, 1), random.uniform(-1, 1), 1))
    emitter = bpy.context.active_object
    emitter.name = "ParticleEmitter"
    
    particle_system = emitter.modifiers.new(name="ParticleSettings", type='PARTICLE_SYSTEM')
    particle_system.particle_system.settings.count = 30
    particle_system.particle_system.settings.frame_start = 1
    particle_system.particle_system.settings.frame_end = duration * bpy.context.scene.render.fps
    particle_system.particle_system.settings.render_type = 'OBJECT'
    particle_system.particle_system.settings.instance_object = bpy.data.objects["ParticleOne"]
    emitter.keyframe_insert(data_path="location", index=-1, frame=1)
    emitter.location = (random.uniform(-10, 10), random.uniform(-10, 10), random.uniform(-10, 10))
    emitter.keyframe_insert(data_path="location", index=-1, frame=20)

def create_meteor_with_particle_system():
    ##wip
    bpy.ops.mesh.primitive_ico_sphere_add(subdivisions=4, radius=0.25, location=(random.uniform(-10, 10), random.uniform(-10, 10), random.uniform(-10, 10)))
    meteor = bpy.context.active_object
    meteor.name = "Meteor"
    meteor_material = bpy.data.materials.new(name="MeteorMaterial")
    meteor.data.materials.append(meteor_material)
    meteor.active_material.use_nodes = True
    nodes = meteor.active_material.node_tree.nodes
    emission_node = nodes.new(type='ShaderNodeEmission')
    emission_node.inputs[0].default_value = (0.0, 0.0, 1.0, 1)  # Blue color
    output_node = nodes.get("Material Output")
    material_link = meteor.active_material.node_tree.links.new
    material_link(emission_node.outputs[0], output_node.inputs[0])
    meteor.location = (random.uniform(-10, 10), random.uniform(-10, 10), random.uniform(-10, 10))
    bpy.context.view_layer.objects.active = meteor
    meteor.select_set(True)
    meteor.animation_data_create()
    meteor.animation_data.action = bpy.data.actions.new(name="OrbitAction")
    start_frame = 1
    duration=2
    radius=3
    end_frame = int(duration * 10)
    for frame in range(start_frame, end_frame + 1):
        angle = (frame / end_frame) * 2 * math.pi
        x = radius * math.cos(angle)
        y = radius * math.sin(angle)
        z = meteor.location.z
        meteor.location = (x, y, z)
        meteor.keyframe_insert(data_path="location", index=-1, frame=frame)

    ###meteor.keyframe_insert(data_path="location", frame=1)
    ##meteor.location = (random.uniform(-10, 10), random.uniform(-10, 10), random.uniform(-10, 10))
    ##meteor.keyframe_insert(data_path="location", frame=20)

def create_scene():

    world = bpy.context.scene.world
    if not world.use_nodes:
        world.use_nodes = True
    ##world = bpy.context.scene.world
    ##bg_node = world.node_tree.nodes['Background']
    ##bg_node.image = img

    ##env_data = bpy.data.worlds['World'].node_tree.nodes['Background'].input
    ##env_data.default_value = img
    bg_node = world.node_tree.nodes.get('Background')
    bg_node.inputs['Color'].default_value = (0, 0, 0, 1)
    
    ##tex_environment = world.node_tree.nodes.new('ShaderNodeTexEnvironment')
    ##tex_environment.image = img
    ##world = bpy.context.scene.world
    ##tx_node = world.node_tree.nodes['Background']
    ##world.node_tree.links.new(tex_environment.outputs[0], tx_node.inputs[0])

    ##img = bpy.data.images.load(image_path)

    bpy.data.objects['Cube'].select_set(True)
    bpy.ops.object.delete()
    bpy.data.objects['Camera'].select_set(True)
    bpy.ops.object.delete()
    
    light_data = bpy.data.lights.new(name="NewLight", type='SUN')
    light_data.energy = 30
    light_object = bpy.data.objects.new(name="NewLight", object_data=light_data)
    scene.collection.objects.link(light_object)
    light_object.location = (3.0, 5.0, 11.0)

    camera_data = bpy.data.cameras.new(name='Camera')
    camera = bpy.data.objects.new('Camera', camera_data)
    bpy.context.collection.objects.link(camera)
    bpy.context.scene.camera = camera
    camera.location = (0, 0, 10)
    camera.rotation_euler = (0, 0, 0)
    ##planet1 = add_planet("earth", 2.5, 1, 0)
    ##planet2 = add_planet("mars", 1.75, 0.5, 1.5)
    add_planet("earth", 2.5, 1, 0)
    add_planet("mars", 1.75, 0.5, 1.5)
    for i in range(30):
        add_animated_space_rocks("spacerock")
    create_meteor_with_particle_system()
    ##create_particle_effects_animation()
    ##create_planet_flyby_animation(planet1, 4, 2)
    ##visualize_planets()
scene = bpy.context.scene
render = scene.render
render.resolution_x = 1920
render.resolution_y = 1080
render.engine = 'BLENDER_EEVEE'
create_scene()

def export_blend():
    bpy.ops.wm.save_as_mainfile(filepath='./exp_scene.blend')
    last_created_time = 0
    last_created_file = None
    apol=[]

    for root, dirs, files in os.walk('.'):
        for file in files:
            file_path = os.path.join(root, file)
            created_time = os.path.getctime(file_path)
            if created_time > last_created_time:
                last_created_time = created_time
                last_created_file = file_path

    apol.append(last_created_file)
    return last_created_file

def render_scene(sny,progress=gr.Progress(track_tqdm=True)):
    bpy.context.scene.render.filepath = './rendered_scene.png'
    bpy.ops.render.render(animation=False, write_still=True)
    last_created_time = 0
    last_created_file = None
    apol=[]

    for root, dirs, files in os.walk('.'):
        for file in files:
            file_path = os.path.join(root, file)
            created_time = os.path.getctime(file_path)
            if created_time > last_created_time:
                last_created_time = created_time
                last_created_file = file_path

    imoge = load_image(last_created_file).convert('RGB')
    apol.append(imoge)
    moog = gr.Model3D(value=sny)
    return apol, sny, moog

def export_glb():
    bpy.ops.export_scene.gltf(filepath='./exp_scene.glb', export_format='GLB')##, export_image_format='AUTO', export_materials='EXPORT', export_cameras=True, export_lights=True)
    last_created_time = 0
    last_created_file = None
    apol=[]

    for root, dirs, files in os.walk('.'):
        for file in files:
            file_path = os.path.join(root, file)
            created_time = os.path.getctime(file_path)
            if created_time > last_created_time:
                last_created_time = created_time
                last_created_file = file_path

    apol.append(last_created_file)
    return last_created_file

def render_anim(sny,progress=gr.Progress(track_tqdm=True)):
    bpy.context.scene.render.image_settings.file_format = 'FFMPEG'
    bpy.context.scene.render.ffmpeg.format = 'MPEG4'
    bpy.context.scene.render.ffmpeg.codec = 'H264'
    bpy.context.scene.render.ffmpeg.constant_rate_factor = 'HIGH'
    bpy.context.scene.frame_start = 1
    bpy.context.scene.frame_end = 20
    bpy.context.scene.render.filepath = './animation.mp4'
    bpy.ops.render.render(animation=True)
    
    last_created_time = 0
    last_created_file = None
    apol=[]

    for root, dirs, files in os.walk('.'):
        for file in files:
            file_path = os.path.join(root, file)
            created_time = os.path.getctime(file_path)
            if created_time > last_created_time:
                last_created_time = created_time
                last_created_file = file_path

    apol.append(last_created_file)
    moog = gr.Model3D(value=sny)
    return apol[0], sny, moog

with gr.Blocks() as iface:
    with gr.Column():
        moog = gr.Model3D()
        ony = gr.Button("render Image")
        twy = gr.Button("render Animation")
        rani = gr.Video()
        rdr = gr.Gallery(columns=1)
        sny = gr.File(value=export_glb(),interactive=False)
        siy = gr.File(value=export_blend(),interactive=False)

    ony.click(fn=render_scene, inputs=[sny], outputs=[rdr,sny,moog])
    twy.click(fn=render_anim, inputs=[sny], outputs=[rani,sny,moog])

iface.queue(max_size=1,api_open=False)
iface.launch(max_threads=20,inline=False,show_api=False)