"""
HF Spaces generate module — called via multiprocessing.Process from app.py.
Waits for --signal-path before touching CUDA, writes incremental JSON to --result-path.
Also usable as a CLI script.
"""

import argparse
import json
import os
import sys
import time
from pathlib import Path

WAN2GP_ROOT = Path(os.environ.get("WAN2GP_ROOT", "/tmp/Wan2GP"))

MODEL_SHORTHANDS = {
    "sulphur-2": "sulphur_2_base",
}

DEFAULTS = {
    "sulphur_2_base": {
        "num_inference_steps": 8,
        "guidance_scale":      5.0,
        "resolution":          "832x480",
        "video_length":        81,
    },
}

_log_entries = []
_result_path = None


def _p(msg):
    ts = time.strftime("%H:%M:%S")
    entry = f"[{ts}] {msg}"
    _log_entries.append(entry)
    print(entry, flush=True)
    if _result_path:
        try:
            with open(_result_path, "w") as f:
                json.dump({"log": _log_entries, "done": False}, f)
        except Exception:
            pass


def _done(error=None):
    if _result_path:
        with open(_result_path, "w") as f:
            data = {"log": _log_entries, "done": True}
            if error:
                data["error"] = error
            json.dump(data, f)


def run_generation(image, prompt, output, model="sulphur-2", steps=None,
                   guidance_scale=None, frames=None, resolution=None, seed=-1,
                   signal_path=None, result_path=None):
    global _log_entries, _result_path
    _log_entries = []
    _result_path = result_path

    if signal_path:
        _p("Waiting for GPU lease signal...")
        while not Path(signal_path).exists():
            time.sleep(0.05)
        _p("Signal received — initialising CUDA context immediately...")
        try:
            import torch
            _ = torch.zeros(1, device="cuda")
            _p(f"CUDA ready: {torch.cuda.get_device_name(0)}")
        except Exception as exc:
            _p(f"CUDA init failed: {exc}")
            _done(error=str(exc))
            return

    model_type = MODEL_SHORTHANDS.get(model, model)
    defaults   = DEFAULTS.get(model_type, DEFAULTS["sulphur_2_base"])

    image_path = str(Path(image.strip()).resolve())
    if not Path(image_path).exists():
        _p(f"Fatal: image not found: {image_path}")
        _done(error="image not found")
        return

    res = resolution or defaults["resolution"]
    if not resolution:
        try:
            from PIL import Image as _PIL
            img = _PIL.open(image_path)
            iw, ih = img.size
            if ih > iw:
                tw = 480; th = round(ih / iw * tw / 32) * 32
            else:
                th = 480; tw = round(iw / ih * th / 32) * 32
            res = f"{tw}x{th}"
            _p(f"Auto-detected resolution: {res} (from {iw}x{ih} input)")
        except Exception:
            pass

    task = {
        "model_type":           model_type,
        "base_model_type":      model_type,
        "prompt":               prompt,
        "image_start":          image_path,
        "num_inference_steps":  steps         or defaults["num_inference_steps"],
        "guidance_scale":       guidance_scale or defaults["guidance_scale"],
        "resolution":           res,
        "video_length":         frames        or defaults["video_length"],
        "seed":                 seed,
        "image_prompt_type":    "S",
        "input_video_strength": 1.0,
        "activated_loras": [
            "ltx-2.3-22b-distilled-lora-384.safetensors",
            "sulphur_experimental_lora_v1.safetensors",
        ],
        "loras_multipliers": ["0.5", "1.0"],
    }

    _p(f"Model:      {model_type}")
    _p(f"Image:      {image_path}")
    _p(f"Steps:      {task['num_inference_steps']}  Guidance: {task['guidance_scale']}")
    _p(f"Resolution: {task['resolution']}  Frames: {task['video_length']}")
    _p(f"Prompt:     {prompt[:80]}")

    sys.path.insert(0, str(WAN2GP_ROOT))
    os.chdir(WAN2GP_ROOT)

    from shared.api import WanGPSession

    output_dir = Path(output).parent
    output_dir.mkdir(parents=True, exist_ok=True)

    _p("Starting WanGPSession...")
    session = WanGPSession(root=WAN2GP_ROOT, output_dir=output_dir, console_output=True)

    _p("Running generation...")
    result = session.run_task(task)

    output_file = None
    if result.artifacts:
        src = result.artifacts[0].path
        if src and Path(src).exists():
            output_file = src

    if output_file is None:
        candidates = sorted(output_dir.glob("**/*.mp4"), key=lambda f: f.stat().st_mtime, reverse=True)
        if candidates:
            output_file = str(candidates[0])
            _p(f"Found output via dir scan: {output_file}")

    if output_file:
        import shutil
        shutil.copy2(output_file, output)
        _p(f"Done: {output}")
    else:
        _p(f"No output found in {output_dir}")
        if result.errors:
            _p(f"Errors: {result.errors}")
        _done(error="no output produced")
        return

    session.close()
    _done()


# Top-level function required for multiprocessing.Process pickling
def _worker_entrypoint(image, prompt, output, model, steps, guidance_scale,
                       frames, resolution, seed, signal_path, result_path):
    run_generation(
        image=image, prompt=prompt, output=output, model=model,
        steps=steps, guidance_scale=guidance_scale, frames=frames,
        resolution=resolution, seed=seed,
        signal_path=signal_path, result_path=result_path,
    )


def parse_args():
    ap = argparse.ArgumentParser()
    ap.add_argument("--image",          required=True)
    ap.add_argument("--prompt",         required=True)
    ap.add_argument("--output",         required=True)
    ap.add_argument("--model",          default="sulphur-2")
    ap.add_argument("--steps",          type=int,   default=None)
    ap.add_argument("--guidance_scale", type=float, default=None)
    ap.add_argument("--frames",         type=int,   default=None)
    ap.add_argument("--resolution",     default=None)
    ap.add_argument("--seed",           type=int,   default=-1)
    ap.add_argument("--signal-path",    default=None, dest="signal_path")
    ap.add_argument("--result-path",    default=None, dest="result_path")
    return ap.parse_args()


def main():
    args = parse_args()
    run_generation(
        image=args.image, prompt=args.prompt, output=args.output,
        model=args.model, steps=args.steps, guidance_scale=args.guidance_scale,
        frames=args.frames, resolution=args.resolution, seed=args.seed,
        signal_path=args.signal_path, result_path=args.result_path,
    )


if __name__ == "__main__":
    main()