#!/usr/bin/env python3
"""
hyworld_gen.py — Generate 3D worlds from text / image / video using HunyuanWorld 2.0.

Integrates Tencent's HY-World-2.0 (https://github.com/Tencent-Hunyuan/HY-World-2.0)
into zelpi's world generation pipeline.

Pipeline (multi-modal, same 3 input modes as zelpi intent):
  text:   text → HY-Pano 2.0 (text→pano) → WorldMirror 2.0 → 3DGS + depth + poses
  image:  image → HY-Pano 2.0 (image→pano) → WorldMirror 2.0 → same outputs
  video:  video → WorldMirror 2.0 directly → 3DGS + depth + camera poses

Output:
  panorama.png       — 360° equirectangular expansion (1952×960)
  points.ply         — colored point cloud (up to 2M points)
  gaussians.ply      — 3D Gaussian Splats (up to 5M Gaussians)
  camera_params.json — camera intrinsics + extrinsics (per-view)
  depth_maps/        — per-view depth visualizations (optional)
  result.json        — metadata and result summary

Usage:
  python scripts/hyworld_gen.py --mode text --input "<prompt>" --out <dir>
  python scripts/hyworld_gen.py --mode image --input image.png --out <dir>
  python scripts/hyworld_gen.py --mode video --input video.mp4 --out <dir> --quality high
"""
from __future__ import annotations

import argparse
import json
import logging
import os
import pathlib
import sys
import time
from typing import Optional

log = logging.getLogger("hyworld_gen")


def _resolve_wm_output_dir(pipeline_result, requested_out: pathlib.Path) -> pathlib.Path:
    """
    WorldMirrorPipeline.__call__() creates its OWN timestamped subdirectory under
    output_path (e.g. <output_path>/<case-name>/<timestamp>/) rather than writing
    flat into output_path, and returns that actual directory as a string. Use it
    when available instead of assuming files land directly in requested_out.
    """
    if isinstance(pipeline_result, (str, pathlib.Path)):
        p = pathlib.Path(pipeline_result)
        if p.exists():
            return p
    return requested_out


def _wm_result_metadata(wm_dir: pathlib.Path, quality: str) -> dict:
    """Build the result dict from wherever WorldMirror actually wrote its output."""
    depth_dir = wm_dir / "depth"
    return {
        "pointcloud": str(wm_dir / "points.ply") if (wm_dir / "points.ply").exists() else None,
        "gaussians": str(wm_dir / "gaussians.ply") if (wm_dir / "gaussians.ply").exists() else None,
        "depth_maps": str(depth_dir) if depth_dir.exists() else None,
        "camera_params": str(wm_dir / "camera_params.json") if (wm_dir / "camera_params.json").exists() else None,
        "quality": quality,
    }


def _find_hyworld_path() -> Optional[pathlib.Path]:
    """Find the HY-World-2.0 repo via env var or common locations."""
    if hyw := os.environ.get("HY_WORLD_PATH"):
        p = pathlib.Path(hyw)
        if (p / "hyworld2").exists():
            return p
    # fallback: check a few common clones
    for base in [pathlib.Path.home(), pathlib.Path.cwd().parent, pathlib.Path.cwd()]:
        for name in ["HY-World-2.0", "hy-world-2.0", "hyworld"]:
            p = base / name
            if (p / "hyworld2").exists():
                return p
    return None


def _add_hyworld_paths(hy_path: pathlib.Path) -> None:
    """Put the repo root AND its script-style subpackage dirs on sys.path.

    HY-World-2.0's pano pipeline (hyworld2/panogen/pipeline_with_qwen_image.py)
    does a BARE `from qwen_image import ...` — qwen_image is vendored as a
    sibling directory, so that import only resolves when hyworld2/panogen
    itself is on sys.path (upstream runs the file as a script from inside
    that directory). Importing it as a package module from here without this
    fails with "No module named 'qwen_image'".
    """
    for p in (hy_path, hy_path / "hyworld2" / "panogen"):
        s = str(p)
        if p.is_dir() and s not in sys.path:
            sys.path.insert(0, s)


def _validate_hyworld(hy_path: Optional[pathlib.Path] = None) -> bool:
    """Check if HY-World-2.0 is importable (optionally with explicit path)."""
    if hy_path:
        _add_hyworld_paths(hy_path)
    try:
        import hyworld2  # noqa: F401
        return True
    except ImportError:
        return False


def _gpu_vram_gb() -> Optional[float]:
    """Return total VRAM of GPU 0 in GB, or None if no CUDA GPU is visible."""
    try:
        import torch
        if torch.cuda.is_available():
            return torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)
    except Exception:
        pass
    return None


def _apply_low_vram(pipeline, name: str, force: bool = False) -> None:
    """
    Best-effort: enable CPU offload / memory-saving modes on a diffusers-style
    pipeline so large models (e.g. the ~20B Qwen-Image-Edit-2509 backbone behind
    HY-Pano-2.0) can still run — slowly — on small GPUs instead of OOMing outright.
    Every call is duck-typed and wrapped, since not all pipelines implement all
    of these methods; skip anything that doesn't apply.
    """
    vram = _gpu_vram_gb()
    if not force and (vram is None or vram >= 16):
        return  # plenty of VRAM (or no GPU info) — don't slow things down needlessly

    log.info("LOW VRAM (%s): enabling CPU offload for %s — this will be slower",
              f"{vram:.1f} GB" if vram else "forced", name)

    for method, label in [
        ("enable_sequential_cpu_offload", "sequential CPU offload"),
        ("enable_model_cpu_offload", "model CPU offload"),
        ("enable_vae_slicing", "VAE slicing"),
        ("enable_vae_tiling", "VAE tiling"),
        ("enable_attention_slicing", "attention slicing"),
    ]:
        fn = getattr(pipeline, method, None)
        if callable(fn):
            try:
                fn()
                log.info("  ✓ %s enabled on %s", label, name)
                if method in ("enable_sequential_cpu_offload", "enable_model_cpu_offload"):
                    break  # these two are mutually exclusive / redundant — one is enough
            except Exception as e:
                log.warning("  ✗ %s failed on %s: %s", label, name, e)


def _validate_input(mode: str, input_path: str) -> bool:
    """Validate that input exists (for image/video modes)."""
    if mode == "text":
        return bool(input_path.strip())
    p = pathlib.Path(input_path)
    if not p.exists():
        log.error("Input file not found: %s", input_path)
        return False
    if mode == "image" and p.suffix.lower() not in {".png", ".jpg", ".jpeg", ".webp"}:
        log.error("Image must be .png, .jpg, .jpeg, or .webp, got %s", p.suffix)
        return False
    # "video" mode also accepts a directory of extracted frames — WorldMirror-2.0
    # natively supports either a video file or an image directory as input.
    if mode == "video" and p.is_dir():
        return True
    if mode == "video" and p.suffix.lower() not in {".mp4", ".avi", ".mov", ".mkv"}:
        log.error("Video must be .mp4, .avi, .mov, or .mkv, got %s", p.suffix)
        return False
    return True


def generate_text_to_3d(prompt: str, out: pathlib.Path, quality: str) -> dict:
    """Text → HY-Pano 2.0 → WorldMirror 2.0 → 3D world."""
    try:
        from hyworld2.panogen.pipeline_with_qwen_image import HunyuanPanoPipeline
        from hyworld2.worldrecon.pipeline import WorldMirrorPipeline
    except ImportError as e:
        return {"error": f"Failed to import HunyuanWorld models: {e}"}

    log.info("TEXT-TO-PANO: generating 360° panorama from prompt")
    pano_pipeline = HunyuanPanoPipeline.from_pretrained(
        pretrained_model_name_or_path="Qwen/Qwen-Image-Edit-2509",
        lora_path="tencent/HY-World-2.0",
        lora_subfolder="HY-Pano-2.0",
    )
    _apply_low_vram(pano_pipeline, "HY-Pano-2.0")

    # Generate panorama from text prompt (HY-Pano 2.0 Backend 2, lightweight)
    try:
        pano_output = pano_pipeline(
            image=None,  # start from text (no input image to edit)
            prompt=prompt,
            seed=42,
            height=960,
            width=1952,
            num_inference_steps=40 if quality == "draft" else 50,
            guidance_scale=1.0,
            blend_width=32,
        )
        pano_path = out / "panorama.png"
        pano_output.save(str(pano_path))
        log.info("PANORAMA saved: %s", pano_path)
    except Exception as e:
        log.error("Panorama generation failed: %s", e)
        return {"error": f"Panorama generation failed: {e}"}

    # Feed panorama into WorldMirror 2.0 for 3D reconstruction
    log.info("PANO-TO-3D: reconstructing 3D world from panorama")
    disable_heads = []
    if quality == "draft":
        disable_heads = ["normal", "camera", "gs"]
    elif quality == "standard":
        disable_heads = ["normal"]

    # WorldMirror expects a directory of images, so wrap the panorama in one
    pano_dir = out / "pano_input"
    pano_dir.mkdir(exist_ok=True)
    import shutil
    shutil.copy(str(pano_path), str(pano_dir / "panorama.png"))

    try:
        wm_pipeline = WorldMirrorPipeline.from_pretrained(
            pretrained_model_name_or_path="tencent/HY-World-2.0",
            subfolder="HY-WorldMirror-2.0",
            use_fsdp=False,
            enable_bf16=False,
            disable_heads=disable_heads,
        )
        _apply_low_vram(wm_pipeline, "WorldMirror-2.0")
        result = wm_pipeline(
            input_path=str(pano_dir),
            output_path=str(out),
            target_size=952,
            save_depth=True,
            save_normal=(quality == "high"),
            save_gs=(quality != "draft"),
            save_camera=True,
            save_points=True,
            compress_pts=(quality == "draft"),
            compress_pts_max_points=1_000_000 if quality == "draft" else 2_000_000,
        )
        log.info("3D RECONSTRUCTION complete: depth, points, gs, camera poses")
    except Exception as e:
        log.error("3D reconstruction failed: %s", e)
        return {"error": f"3D reconstruction failed: {e}"}

    wm_dir = _resolve_wm_output_dir(result, out)
    return {"panorama": str(pano_path), **_wm_result_metadata(wm_dir, quality)}


def generate_image_to_3d(image_path: str, out: pathlib.Path, quality: str) -> dict:
    """Image → HY-Pano 2.0 (image expansion) → WorldMirror 2.0 → 3D world."""
    try:
        from hyworld2.panogen.pipeline_with_qwen_image import HunyuanPanoPipeline
        from hyworld2.worldrecon.pipeline import WorldMirrorPipeline
        from PIL import Image
    except ImportError as e:
        return {"error": f"Failed to import HunyuanWorld models: {e}"}

    log.info("IMAGE-TO-PANO: expanding perspective image to 360° panorama")
    img = Image.open(image_path).convert("RGB")

    pano_pipeline = HunyuanPanoPipeline.from_pretrained(
        pretrained_model_name_or_path="Qwen/Qwen-Image-Edit-2509",
        lora_path="tencent/HY-World-2.0",
        lora_subfolder="HY-Pano-2.0",
    )
    _apply_low_vram(pano_pipeline, "HY-Pano-2.0")

    # Expand image to panorama
    try:
        pano_output = pano_pipeline(
            image=img,
            prompt="Expand this image to a 360-degree equirectangular panorama.",
            seed=42,
            height=960,
            width=1952,
            num_inference_steps=40 if quality == "draft" else 50,
            guidance_scale=1.0,
            blend_width=32,
        )
        pano_path = out / "panorama.png"
        pano_output.save(str(pano_path))
        log.info("PANORAMA saved: %s", pano_path)
    except Exception as e:
        log.error("Panorama generation failed: %s", e)
        return {"error": f"Panorama generation failed: {e}"}

    # Feed panorama into WorldMirror 2.0 for 3D reconstruction
    log.info("PANO-TO-3D: reconstructing 3D world from panorama")
    disable_heads = []
    if quality == "draft":
        disable_heads = ["normal", "camera", "gs"]
    elif quality == "standard":
        disable_heads = ["normal"]

    # WorldMirror expects a directory of images, so wrap the panorama in one
    pano_dir = out / "pano_input"
    pano_dir.mkdir(exist_ok=True)
    import shutil
    shutil.copy(str(pano_path), str(pano_dir / "panorama.png"))

    try:
        wm_pipeline = WorldMirrorPipeline.from_pretrained(
            pretrained_model_name_or_path="tencent/HY-World-2.0",
            subfolder="HY-WorldMirror-2.0",
            use_fsdp=False,
            enable_bf16=False,
            disable_heads=disable_heads,
        )
        _apply_low_vram(wm_pipeline, "WorldMirror-2.0")
        result = wm_pipeline(
            input_path=str(pano_dir),
            output_path=str(out),
            target_size=952,
            save_depth=True,
            save_normal=(quality == "high"),
            save_gs=(quality != "draft"),
            save_camera=True,
            save_points=True,
            compress_pts=(quality == "draft"),
            compress_pts_max_points=1_000_000 if quality == "draft" else 2_000_000,
        )
        log.info("3D RECONSTRUCTION complete")
    except Exception as e:
        log.error("3D reconstruction failed: %s", e)
        return {"error": f"3D reconstruction failed: {e}"}

    wm_dir = _resolve_wm_output_dir(result, out)
    return {"panorama": str(pano_path), **_wm_result_metadata(wm_dir, quality)}


def generate_video_to_3d(video_path: str, out: pathlib.Path, quality: str) -> dict:
    """Video / image sequence → WorldMirror 2.0 → 3D world (depth + poses + GS)."""
    try:
        from hyworld2.worldrecon.pipeline import WorldMirrorPipeline
    except ImportError as e:
        return {"error": f"Failed to import HunyuanWorld models: {e}"}

    log.info("VIDEO-TO-3D: reconstructing 3D world from video / image sequence")
    disable_heads = []
    if quality == "draft":
        disable_heads = ["normal", "gs"]
    elif quality == "standard":
        disable_heads = ["normal"]

    try:
        wm_pipeline = WorldMirrorPipeline.from_pretrained(
            pretrained_model_name_or_path="tencent/HY-World-2.0",
            subfolder="HY-WorldMirror-2.0",
            use_fsdp=False,
            enable_bf16=False,
            disable_heads=disable_heads,
        )
        _apply_low_vram(wm_pipeline, "WorldMirror-2.0")
        result = wm_pipeline(
            input_path=video_path,
            output_path=str(out),
            target_size=952,
            fps=2,  # extract 2 frames per second for dense reconstruction
            save_depth=True,
            save_normal=(quality == "high"),
            save_gs=(quality != "draft"),
            save_camera=True,
            save_points=True,
            compress_pts=(quality == "draft"),
            compress_pts_max_points=1_000_000 if quality == "draft" else 2_000_000,
        )
        log.info("3D RECONSTRUCTION complete: depth, points, gs, camera poses")
    except Exception as e:
        log.error("3D reconstruction failed: %s", e)
        return {"error": f"3D reconstruction failed: {e}"}

    wm_dir = _resolve_wm_output_dir(result, out)
    return _wm_result_metadata(wm_dir, quality)


def main():
    ap = argparse.ArgumentParser(description="Generate 3D worlds using HunyuanWorld 2.0")
    ap.add_argument("--mode", choices=["text", "image", "video"], required=True)
    ap.add_argument("--input", required=True, help="text prompt, image path, or video path")
    ap.add_argument("--out", required=True, help="output directory")
    ap.add_argument("--quality", choices=["draft", "standard", "high"], default="standard")
    args = ap.parse_args()

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s  %(name)s  %(levelname)s  %(message)s",
        datefmt="%H:%M:%S",
    )

    # Check HY-World-2.0 installation
    hyworld_path = _find_hyworld_path()
    if not _validate_hyworld(hyworld_path):
        result = {
            "code": "not_installed",
            "message": "HY-World-2.0 not found. Install and set HY_WORLD_PATH environment variable.",
            "hyworld_path": str(hyworld_path) if hyworld_path else None,
        }
        print("HYWORLD_RESULT " + json.dumps(result))
        sys.exit(1)
    # Ensure the repo root + script-style subpackage dirs are importable
    # (see _add_hyworld_paths for the vendored-qwen_image trap).
    if hyworld_path:
        _add_hyworld_paths(hyworld_path)

    # Validate input
    if not _validate_input(args.mode, args.input):
        result = {"code": "invalid_input", "message": f"Invalid input for mode {args.mode}"}
        print("HYWORLD_RESULT " + json.dumps(result))
        sys.exit(1)

    # Prepare output directory
    out = pathlib.Path(args.out)
    out.mkdir(parents=True, exist_ok=True)

    log.info("HYWORLD %s mode, quality=%s, input=%s", args.mode, args.quality, args.input)

    # Generate 3D world based on mode
    result = None
    if args.mode == "text":
        result = generate_text_to_3d(args.input, out, args.quality)
    elif args.mode == "image":
        result = generate_image_to_3d(args.input, out, args.quality)
    elif args.mode == "video":
        result = generate_video_to_3d(args.input, out, args.quality)

    # Handle errors
    if result and "error" in result:
        log.error("Generation failed: %s", result["error"])
        result["code"] = "generation_failed"
        print("HYWORLD_RESULT " + json.dumps(result))
        sys.exit(1)

    # Success: add metadata and print result sentinel
    result["out_dir"] = str(out)
    result["mode"] = args.mode
    result["generated"] = time.strftime("%Y-%m-%dT%H:%M:%S")

    (out / "result.json").write_text(json.dumps(result, indent=2))
    log.info("HYWORLD generation complete → %s", out)
    print("HYWORLD_RESULT " + json.dumps(result))


if __name__ == "__main__":
    main()