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.initialized = False
        self._setup_blender()
        self._setup_infinigen()

    def _setup_blender(self):
        """Stellt sicher, dass Blender installiert ist und findet den Python-Interpreter."""
        try:
            blender_bins = glob.glob("/usr/bin/blender")
            if not blender_bins:
                raise FileNotFoundError("Blender-Binär nicht gefunden!")
            
            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'}")
            
            # Breite Suche nach Python-Interpreter
            base_paths = ["/usr/share/blender", "/usr/lib/blender"]
            python_patterns = [
                f"/usr/share/blender/{blender_version}/python/bin/python*",
                f"/usr/share/blender/*/python/bin/python*",
                f"/usr/lib/blender/*/python/bin/python*",
                "/**/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}")
                    self.initialized = True
                    return
            
            # Debugging: Vollständige Python-Suche
            all_python = glob.glob("/**/python", recursive=True)
            print(f"Alle Python-Pfade: {all_python}")
            raise RuntimeError("Blender Python-Interpreter nicht gefunden!")
        except Exception as e:
            print(f"Fehler bei Blender-Setup: {e}")
            self.blender_python = None

    def _setup_infinigen(self):
        """Installiert Infinigen, falls möglich."""
        if not self.blender_python:
            print("Kann Infinigen nicht installieren: Blender Python fehlt.")
            return
        
        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."""
        if not self.blender_python:
            return False
        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."""
        if not self.initialized or not self.blender_python:
            return "Fehler: Infinigen nicht initialisiert (Blender Python fehlt)!"
        
        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
manager = InfinigenManager()  # Immer definieren, auch bei Fehlern

# 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()