from __future__ import annotations

import secrets
import json
import threading
import time
from datetime import datetime, timezone
from typing import Callable

import torch
from PIL import Image

from .model_resolver import prepare_models_root
from .realesrgan_upscaler import ResidentHDUpscaler
from .space_config import SpaceConfig
from .space_postprocess import (
    apply_rife_seam,
    close_rife_model,
    prepare_rife_model,
    save_frame_bundle,
)
from .wanvideo_loop_runtime import WrapperLoopRuntime


def _diag(event: str, **fields) -> None:
    payload = {
        "ts": datetime.now(timezone.utc).isoformat(timespec="milliseconds"),
        "event": event,
        **fields,
    }
    print(f"[WAN_SERVICE] {json.dumps(payload, sort_keys=True)}", flush=True)


class LoopGeneratorService:
    def __init__(self, config: SpaceConfig) -> None:
        self.config = config
        started = time.perf_counter()
        _diag("service.init.start")
        models_root, names = prepare_models_root()
        self.runtime = WrapperLoopRuntime(
            models_root=models_root,
            high_model_name=names["high"],
            low_model_name=names["low"],
            clip_name=names["text"],
            int8_clip_name=names["text_int8"],
            vae_name=names["vae"],
            sampler_name=config.sampler,
            scheduler_mode=config.scheduler,
            split_step=config.split_step,
            riflex_k=0,
            loop_shift_skip=0,
            loop_start_percent=0.0,
            loop_end_percent=1.0,
            start_latent_strength=config.start_strength,
            end_latent_strength=config.end_strength,
            end_temporal_mask_strength=config.end_mask_strength,
            decode_end_image_hint=True,
            fun_or_fl2v_model=False,
            zero_end_latent_conditioning=False,
            end_latent_conditioning_strength=1.0,
            low_pass_end_conditioning_strength=1.0,
            custom_sigmas=(),
            attention_mode="sdpa",
            text_encoder_quantization=config.text_encoder_quantization,
            global_resident_models=config.global_resident_models,
            vae_tiling=config.vae_tiling,
        )
        if config.use_rife:
            prepare_rife_model()
        self.upscaler = ResidentHDUpscaler()
        self._job_lock = threading.Lock()
        _diag("service.init.done", elapsed_s=round(time.perf_counter() - started, 3))

    def cleanup_job(self) -> None:
        self.runtime.cleanup_job()

    def close(self) -> None:
        """Destroy process-global resources. Do not call after a normal job."""
        self.upscaler.close()
        close_rife_model()
        self.runtime.shutdown()

    def generate_iter(
        self,
        image: Image.Image,
        prompt: str,
        progress_callback: Callable[..., None] | None = None,
    ):
        if image is None:
            raise ValueError("Upload an image.")
        if not (prompt or "").strip():
            raise ValueError("Enter a prompt.")
        if not self._job_lock.acquire(blocking=False):
            raise RuntimeError("The generator is already processing another request.")
        image = image.convert("RGB")
        frames = None
        try:
            self.runtime.cleanup_job()
            with torch.inference_mode():
                for stage_result in self.runtime.generate_segment_iter(
                    prompt=prompt,
                    negative_prompt="",
                    start_image=image,
                    end_image=image,
                    width=self.config.width,
                    height=self.config.height,
                    num_frames=self.config.frame_count,
                    steps=self.config.steps,
                    cfg=self.config.cfg,
                    shift=self.config.shift,
                    seed=secrets.randbits(63),
                    progress_callback=progress_callback,
                ):
                    if isinstance(stage_result, dict) and "stage" in stage_result:
                        yield stage_result
                    else:
                        frames = stage_result
            if frames is None:
                raise RuntimeError("Frame generation ended without decoded frames.")
            metrics = dict(self.runtime.last_metrics)
            if self.config.use_rife:
                if progress_callback is not None:
                    progress_callback(0.87, desc="Smoothing loop seam…")
                yield {"stage": "Smoothing loop seam…"}
                if torch.cuda.is_available():
                    torch.cuda.synchronize()
                rife_started = time.perf_counter()
                _diag("rife.start", input_frames=int(frames.shape[0]))
                frames = apply_rife_seam(frames, self.config.rife_frames)
                if torch.cuda.is_available():
                    torch.cuda.synchronize()
                metrics["rife_s"] = time.perf_counter() - rife_started
                if torch.cuda.is_available():
                    metrics["peak_vram_gib"] = torch.cuda.max_memory_allocated() / (1024**3)
                _diag(
                    "rife.done",
                    elapsed_s=round(metrics["rife_s"], 3),
                    output_frames=int(frames.shape[0]),
                    output_device=str(frames.device),
                )
            if progress_callback is not None:
                progress_callback(0.91, desc="Upscaling to HD…")
            yield {"stage": "Upscaling to HD…"}
            if torch.cuda.is_available():
                torch.cuda.synchronize()
            upscale_started = time.perf_counter()
            _diag(
                "upscale.start",
                input_width=int(frames.shape[2]),
                input_height=int(frames.shape[1]),
                input_frames=int(frames.shape[0]),
            )
            frames = self.upscaler.upscale_to_hd(frames)
            if torch.cuda.is_available():
                torch.cuda.synchronize()
            upscale_s = time.perf_counter() - upscale_started
            metrics["upscale_s"] = upscale_s
            if torch.cuda.is_available():
                metrics["peak_vram_gib"] = torch.cuda.max_memory_allocated() / (1024**3)
            _diag(
                "upscale.done",
                elapsed_s=round(upscale_s, 3),
                output_width=int(frames.shape[2]),
                output_height=int(frames.shape[1]),
                output_frames=int(frames.shape[0]),
                output_device=str(frames.device),
            )
            print(f"[SPACE_METRICS_FINAL] {json.dumps(metrics, sort_keys=True)}", flush=True)
            if progress_callback is not None:
                progress_callback(0.95, desc="Preparing frames…")
            yield {"stage": "Preparing frames…"}
            bundle_started = time.perf_counter()
            _diag("frame_bundle.start", input_device=str(frames.device))
            bundle = save_frame_bundle(frames)
            _diag("frame_bundle.done", elapsed_s=round(time.perf_counter() - bundle_started, 3))
            yield {"bundle": bundle}
        finally:
            frames = None
            self.runtime.cleanup_job()
            self._job_lock.release()