import os
import shutil
import tempfile
import threading
from pathlib import Path

# === Runtime environment tweaks ===
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")  # force CPU
os.environ.setdefault("WORLDGEN_DISABLE_EVAL", "1")  # avoid KNN/eval paths if supported
os.environ.setdefault("GRADIO_TEMP_DIR", "/tmp/gradio")  # centralize gradio cache

# Monkey patch get_precision to avoid CUDA errors
import nunchaku.utils

def safe_get_precision(*args, **kwargs):
    return "fp32"

nunchaku.utils.get_precision = safe_get_precision

try:
    import nunchaku.models.transformers.transformer_flux as tf
    tf.get_precision = safe_get_precision
except Exception:
    pass

import gradio as gr
from worldgen import WorldGen

_wg = None
def get_worldgen():
    global _wg
    if _wg is None:
        _wg = WorldGen(mode="t2s", device="cpu", low_vram=True, lora_path=None)
    return _wg

def _zip_directory(dir_path: Path) -> str:
    zip_path = dir_path.with_suffix(".zip")
    if zip_path.exists():
        zip_path.unlink()
    shutil.make_archive(str(dir_path), "zip", root_dir=str(dir_path))
    return str(zip_path)

def _save_scene(generated, out_dir: Path) -> Path:
    out_dir.mkdir(parents=True, exist_ok=True)

    for meth in ("save", "export", "write", "to_disk"):
        if hasattr(generated, meth):
            getattr(generated, meth)(str(out_dir))
            return out_dir

    for meth in ("save_as_glb", "to_glb", "save_as_ply", "to_ply"):
        if hasattr(generated, meth):
            fp = out_dir / ("scene.glb" if "glb" in meth else "scene.ply")
            getattr(generated, meth)(str(fp))
            return out_dir

    (out_dir / "README.txt").write_text(
        "No known save/export method detected. Please adjust _save_scene according to your WorldGen version."
    )
    return out_dir

def generate_from_text(prompt: str):
    if not prompt or not prompt.strip():
        return None, "Please enter a scene description (English works best)."

    wg = get_worldgen()

    try:
        result = wg.generate_world(prompt.strip())
    except Exception as e:
        return None, f"Generation failed: {e}"

    tmpdir = Path(tempfile.mkdtemp())
    out_dir = tmpdir / "worldgen_output"
    out_dir = _save_scene(result, out_dir)
    zip_path = _zip_directory(out_dir)

    def _cleanup(p: Path):
        try:
            shutil.rmtree(p, ignore_errors=True)
        except Exception:
            pass

    threading.Timer(300, _cleanup, args=(tmpdir,)).start()

    return zip_path, "Done. Download the ZIP (contains 3D assets)."

with gr.Blocks(title="WorldGen • Text -> 3D Scene") as demo:
    gr.Markdown(
        """# WorldGen - Text -> 3D
- CPU-only fallback is enabled. For real runs, switch the Space to a GPU (T4/A10G).
- Outputs are zipped for download; temp files are cleaned shortly after each run.
"""
    )
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Scene prompt (English)",
                placeholder="a neon-lit cyberpunk street with rain reflections, highly detailed",
                lines=3,
            )
            btn = gr.Button("Generate 3D Scene")
        with gr.Column():
            out_zip = gr.File(label="Download output (ZIP)")
            status = gr.Markdown()

    btn.click(generate_from_text, inputs=[prompt], outputs=[out_zip, status])

demo.launch()