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

class InfinigenManager:
    def __init__(self, base_dir="infinigen", output_dir="outputs"):
        self.base_dir = Path(base_dir)
        self.output_dir = Path(output_dir)
        self.blender_bin = None
        self.blender_python = None
        self._setup_blender()
        self._setup_infinigen()

    def _setup_blender(self):
        """Stellt sicher, dass die neueste Blender-Version installiert ist."""
        try:
            # Prüfe vorhandene Blender-Version
            blender_bins = glob.glob("/usr/bin/blender")
            if blender_bins:
                self.blender_bin = blender_bins[0]
                version_result = subprocess.run([self.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"Vorhandene Blender-Version: {blender_version or 'unbekannt'}")
            else:
                raise FileNotFoundError("Blender nicht vorhanden")
            
            # Finde Blender-Python
            base_path = "/usr/share/blender"
            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_bins = glob.glob(pattern, recursive=True)
                if python_bins:
                    self.blender_python = python_bins[0]
                    print(f"Blender Python gefunden: {self.blender_python}")
                    return
            
            raise RuntimeError("Blender Python-Interpreter nicht gefunden!")
        except Exception as e:
            print(f"Fehler bei Blender-Setup: {e}")
            raise

    def _setup_infinigen(self):
        """Installiert Infinigen, falls nicht vorhanden."""
        if not self.base_dir.exists():
            print("Klone Infinigen...")
            subprocess.run(["git", "clone", "https://github.com/princeton-vl/infinigen.git"], check=True)
        
        if not self._is_infinigen_installed():
            print("Installiere Infinigen...")
            subprocess.run([
                self.blender_python, "-m", "pip", "install", "-e", str(self.base_dir) + "[terrain,vis]", "--no-deps", "--user"
            ], check=True)
            print("Infinigen installiert.")

    def _is_infinigen_installed(self):
        """Prüft, ob Infinigen bereits installiert ist."""
        try:
            subprocess.run([self.blender_python, "-c", "import infinigen"], check=True, capture_output=True)
            return True
        except subprocess.CalledProcessError:
            return False

    @spaces.GPU
    def generate_scene(self, seed, configs=None, pipeline_configs=None):
        """Generiert eine Szene mit Infinigen."""
        self.output_dir.mkdir(exist_ok=True)
        
        configs = configs or ["infinigen_examples/configs/desert.gin", "infinigen_examples/configs/simple.gin"]
        pipeline_configs = pipeline_configs or [
            "infinigen_examples/configs/local_16GB.gin",
            "infinigen_examples/configs/monocular.gin",
            "infinigen_examples/configs/blender_gt.gin"
        ]
        
        command = [
            self.blender_python, "-m", "infinigen.datagen.manage_jobs",
            "--output_folder", str(self.output_dir),
            "--num_scenes", "1",
            "--specific_seed", str(int(seed))
        ] + ["--configs"] + configs + ["--pipeline_configs"] + pipeline_configs
        
        try:
            result = subprocess.run(command, capture_output=True, text=True, check=True)
            print(f"STDOUT: {result.stdout}")
            print(f"STDERR: {result.stderr}")
            output_path = self.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}"

# Manager initialisieren
try:
    manager = InfinigenManager()
except Exception as e:
    print(f"Fehler bei der Initialisierung: {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=manager.generate_scene, inputs=[seed_input], outputs=[output_image])

demo.launch()