import gradio as gr
import subprocess
import os
import spaces
from pathlib import Path

# Funktion zum automatischen Finden des Blender-Python-Interpreters
def find_blender_python():
    # Suche nach Blender-Binär
    blender_bin = next(iter(Path("/").glob("**/blender")), None)
    if not blender_bin:
        raise RuntimeError("Blender-Binär nicht gefunden!")
    
    print(f"Blender-Binär gefunden: {blender_bin}")
    
    # Erkenne Blender-Version
    version_result = subprocess.run([str(blender_bin), "-v"], capture_output=True, text=True)
    blender_version = next((line.split()[1] for line in version_result.stdout.split("\n") if "Blender" in line and "." in line), None)
    print(f"Blender-Version erkannt: {blender_version or 'unbekannt'}")
    
    # Suche nach Python-Interpreter
    base_path = Path("/usr/share/blender") if blender_version else Path("/")
    python_patterns = [
        f"{base_path}/{blender_version}/python/bin/python*",
        f"{base_path}/*/python/bin/python*",
        "/usr/lib/blender/*/python/bin/python*"
    ]
    
    for pattern in python_patterns:
        python_bin = next(iter(Path().glob(pattern)), None)
        if python_bin and python_bin.is_file():
            print(f"Blender Python gefunden: {python_bin}")
            return str(python_bin)
    
    raise RuntimeError("Blender Python-Interpreter nicht gefunden!")

# Funktion zum Installieren von Infinigen
def install_infinigen():
    infinigen_dir = Path("infinigen")
    if not infinigen_dir.exists():
        print("Klone Infinigen...")
        subprocess.run(["git", "clone", "https://github.com/princeton-vl/infinigen.git"], check=True)
        
        blender_python = find_blender_python()
        subprocess.run([
            blender_python, "-m", "pip", "install", "-e", str(infinigen_dir) + "[terrain,vis]", "--no-deps", "--user"
        ], check=True)
        print("Infinigen installiert.")

# Infinigen beim Start installieren
try:
    install_infinigen()
except Exception as e:
    print(f"Fehler bei Infinigen-Installation: {e}")

# Funktion zum Generieren einer Szene mit ZeroGPU
@spaces.GPU
def generate_scene(seed):
    output_dir = Path("outputs")
    output_dir.mkdir(exist_ok=True)
    
    try:
        blender_python = find_blender_python()
        command = [
            blender_python, "-m", "infinigen.datagen.manage_jobs",
            "--output_folder", str(output_dir),
            "--num_scenes", "1",
            "--specific_seed", str(int(seed)),
            "--configs", "infinigen/infinigen_examples/configs/desert.gin",
            "infinigen/infinigen_examples/configs/simple.gin",
            "--pipeline_configs", "infinigen/infinigen_examples/configs/local_16GB.gin",
            "infinigen/infinigen_examples/configs/monocular.gin",
            "infinigen/infinigen_examples/configs/blender_gt.gin"
        ]
        result = subprocess.run(command, capture_output=True, text=True, check=True)
        print(f"STDOUT: {result.stdout}")
        print(f"STDERR: {result.stderr}")
        
        output_path = output_dir / "0000000000.png"
        return str(output_path) if output_path.exists() else f"Fehler: Bild nicht gefunden. STDERR: {result.stderr}"
    except subprocess.CalledProcessError as e:
        return f"Fehler: {e.stderr}"
    except Exception as e:
        return f"Fehler beim Generieren: {e}"

# Gradio-Oberfläche
with gr.Blocks(title="Infinigen Demo") as demo:
    gr.Markdown("## Infinigen Scene Generator")
    seed_input = gr.Number(label="Seed", value=0, precision=0)
    output_image = gr.Image(label="Generierte Szene")
    generate_button = gr.Button("Szene generieren")
    
    generate_button.click(fn=generate_scene, inputs=[seed_input], outputs=[output_image])

demo.launch()