import sys
from pathlib import Path
import tempfile
import subprocess
import torch
import torch.nn.functional as F
import torchaudio
import os
from typing import Any
import time
from contextlib import contextmanager
import gradio as gr
import json
import logging
import socket
import numpy as np
import random
from PIL import Image, ImageOps

APP_ROOT = Path(__file__).resolve().parent
MODEL_DIR = APP_ROOT / "hf_models"
MODEL_DIR.mkdir(parents=True, exist_ok=True)

APP_LOG_LEVEL = os.getenv("APP_LOG_LEVEL", "INFO").upper()
REPLICA_ID = os.getenv("SPACE_REPLICA_ID") or socket.gethostname()

class JsonFormatter(logging.Formatter):
    def format(self, record):
        payload = {
            "ts": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(record.created)),
            "level": record.levelname,
            "logger": record.name,
            "msg": record.getMessage(),
            "replica_id": REPLICA_ID,
        }
        extra_fields = getattr(record, "extra_fields", None)
        if isinstance(extra_fields, dict):
            payload.update(extra_fields)
        if record.exc_info:
            payload["exc_info"] = self.formatException(record.exc_info)
        return json.dumps(payload, ensure_ascii=False)

logger = logging.getLogger("ltx_space")
logger.setLevel(APP_LOG_LEVEL)
logger.propagate = False

handler = logging.StreamHandler(sys.stdout)
handler.setFormatter(JsonFormatter())
logger.handlers.clear()
logger.addHandler(handler)

def log_event(message, **fields):
    logger.info(message, extra={"extra_fields": fields})

def log_error(message, **fields):
    logger.exception(message, extra={"extra_fields": fields})

@contextmanager
def timer(name: str):
    start = time.time()
    log_event(f"{name}...")
    yield
    log_event(f"  -> {name} completed in {time.time() - start:.2f} sec")

current_dir = Path(__file__).parent
sys.path.insert(0, str(current_dir / "packages" / "ltx-pipelines" / "src"))
sys.path.insert(0, str(current_dir / "packages" / "ltx-core" / "src"))

import spaces
from huggingface_hub import hf_hub_download, snapshot_download
from ltx_pipelines.distilled import DistilledPipeline
from ltx_core.model.video_vae import TilingConfig
from ltx_pipelines.utils.constants import (
    DEFAULT_SEED, DEFAULT_1_STAGE_HEIGHT, DEFAULT_1_STAGE_WIDTH
)
from ltx_pipelines.utils import ModelLedger
from ltx_pipelines.utils.helpers import generate_enhanced_prompt
from ltx_core.loader.primitives import LoraPathStrengthAndSDOps
from ltx_core.loader.sd_ops import LTXV_LORA_COMFY_RENAMING_MAP

MAX_SEED = np.iinfo(np.int32).max

MAX_SAFE_FRAMES = 265  # 264 + 1 (multiple of 8 + 1), ~14.7 s @ 18 fps

with timer("Downloading and Loading Models"):
    checkpoint_path = hf_hub_download(repo_id="Lightricks/LTX-2", filename="ltx-2-19b-dev.safetensors", local_dir=MODEL_DIR)
    gemma_path = snapshot_download(repo_id="unsloth/gemma-3-12b-it-qat-bnb-4bit", local_dir=MODEL_DIR)
    upsampler_path = hf_hub_download(repo_id="Lightricks/LTX-2", filename="ltx-2-spatial-upscaler-x2-1.0.safetensors", local_dir=MODEL_DIR)
    lora_path = hf_hub_download(repo_id="Lightricks/LTX-2", filename="ltx-2-19b-distilled-lora-384.safetensors", local_dir=MODEL_DIR)

pipeline = DistilledPipeline(
    device=torch.device("cuda"),
    checkpoint_path=checkpoint_path,
    spatial_upsampler_path=upsampler_path,
    gemma_root=gemma_path,
    loras=[LoraPathStrengthAndSDOps(path=lora_path, strength=0.6, sd_ops=LTXV_LORA_COMFY_RENAMING_MAP)],
    fp8transformer=False,
    local_files_only=False
)

text_encoder = pipeline.model_ledger.text_encoder()
text_encoder._init_image_processor()
pipeline._video_encoder = pipeline.model_ledger.video_encoder()
pipeline._transformer = pipeline.model_ledger.transformer()

def extract_last_frame(video_path):
    import cv2
    cap = cv2.VideoCapture(video_path)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    cap.set(cv2.CAP_PROP_POS_FRAMES, max(0, total_frames - 1))
    ret, frame = cap.read()
    cap.release()
    if ret:
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        return Image.fromarray(frame_rgb)
    return None

GENERATION_FPS = 18.0

def compute_num_frames(duration: float) -> int:
    raw = round((duration * GENERATION_FPS) / 8) * 8 + 1
    raw = max(raw, 9)
    raw = min(raw, MAX_SAFE_FRAMES)
    return raw

def get_tiling_config(num_frames: int, height: int, width: int) -> TilingConfig:
    if num_frames > 176:
        return TilingConfig(
            spatial_tile_height=256,
            spatial_tile_width=256,
            temporal_tile_size=32,
        )
    if num_frames > 97:
        return TilingConfig(
            spatial_tile_height=384,
            spatial_tile_width=384,
            temporal_tile_size=48,
        )
    return TilingConfig.default()

@spaces.GPU(duration=60)
def generate_screenplay_gemma(prompt_text, image_path=None):
    char_desc = ""
    if image_path:
        char_msg = [
            {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "Describe the main character in this image in extreme detail (hair, clothes, age, features) in one short paragraph, so I can recreate them in other scenes. Output only the description and nothing else."}]}
        ]
        img = Image.open(image_path).convert("RGB")
        inputs = text_encoder.processor(text=text_encoder.processor.tokenizer.apply_chat_template(char_msg, tokenize=False, add_generation_prompt=True), images=img, return_tensors="pt").to("cuda")
        with torch.inference_mode():
            outputs = text_encoder.model.generate(**inputs, max_new_tokens=256, do_sample=False)
        char_desc = text_encoder.processor.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()
    system_instructions = (
        "You are a professional screenwriter. Write a coherent, chronologically sequential screenplay divided into exactly 8 distinct scenes based on the user's topic. "
        "Each scene's prompt must start with a transition phrase like 'Continuing from the previous frame, [character description] ...' to maintain smooth sequence continuity. "
        "Each scene must have a detailed visual description (1 paragraph) that can be used directly as a text-to-video generation prompt. "
        "If a character description is provided, you MUST include and adapt that exact character description in every single scene's prompt to ensure absolute physical consistency. "
        "Output ONLY a raw, valid JSON list of objects, where each object has 'scene' (integer, 1 to 8) and 'prompt' (string) fields. "
        "Do NOT write any introduction, markdown, codeblocks, explanations, or trailing text. Just the raw JSON."
    )
    user_content = f"Topic: {prompt_text}"
    if char_desc:
        user_content += f"\n\nCharacter Description to use in all scenes: {char_desc}"
    messages = [
        {"role": "system", "content": system_instructions},
        {"role": "user", "content": user_content}
    ]
    prompt = text_encoder.processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = text_encoder.processor(text=prompt, return_tensors="pt").to("cuda")
    with torch.inference_mode():
        outputs = text_encoder.model.generate(**inputs, max_new_tokens=1024, do_sample=True, temperature=0.7)
    raw_response = text_encoder.processor.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()
    return raw_response, char_desc

@spaces.GPU(duration=180)
def generate_single_scene_gpu(prompt, ref_image_path, num_frames, height, width, seed):
    """
    Generate ONE scene in its own GPU allocation (duration=180 s each).
    Returns (clip_bytes, last_frame_image_or_None) so the caller never
    needs to touch a file after the GPU slot is released.
    """
    import gc
    output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    imgs = []
    if ref_image_path:
        try:
            proc = ImageOps.fit(
                Image.open(ref_image_path), (width, height),
                method=Image.LANCZOS, centering=(0.5, 0.5)
            )
            tmp_img = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
            proc.save(tmp_img)
            imgs.append((tmp_img, 0, 1.0))
        except Exception as e:
            log_error(f"Error cropping reference image: {e}")
            imgs.append((ref_image_path, 0, 1.0))

    tiling = get_tiling_config(num_frames, height, width)
    clip_bytes = None
    last_frame_img = None

    try:
        text_encoder.to("cuda")
        with torch.inference_mode():
            v, a, _ = text_encoder(prompt)
        v_ctx = v.to("cuda", non_blocking=True)
        a_ctx = a.to("cuda", non_blocking=True)
        with torch.inference_mode():
            pipeline(
                prompt=prompt,
                output_path=output_path,
                seed=seed,
                height=height,
                width=width,
                num_frames=num_frames,
                frame_rate=GENERATION_FPS,
                images=imgs,
                video_context=v_ctx,
                audio_context=a_ctx,
                input_waveform=None,
                input_waveform_sample_rate=None,
                tiling_config=tiling,
            )

        # Read bytes and last frame BEFORE cleanup so files still exist
        with open(output_path, "rb") as f:
            clip_bytes = f.read()

        # Extract last frame for continuity while the file is still on disk
        last_frame_img = extract_last_frame(output_path)

    finally:
        # Delete temp video file
        try:
            os.remove(output_path)
        except Exception:
            pass
        # Release VRAM
        gc.collect()
        torch.cuda.empty_cache()
        if torch.cuda.is_available():
            torch.cuda.synchronize()

    return clip_bytes, last_frame_img

def generate_all_scenes(scenes_list, character_image, duration, height, width, seed):
    """
    Orchestrator: calls generate_single_scene_gpu for each scene sequentially.
    Each scene gets its own @spaces.GPU slot → no timeout accumulation,
    VRAM is fully released between scenes.

    Continuity chain:
      Scene 1 → ref = character_image (or None)
      Scene N → ref = PIL Image of the last frame from Scene N-1
                       (returned by generate_single_scene_gpu before cleanup)
    """
    import gc
    clips_bytes = []
    current_ref = character_image   # file path string or None
    num_frames = compute_num_frames(duration)
    log_event(
        f"Generating {len(scenes_list)} scenes | "
        f"requested={duration}s | actual={((num_frames-1)/GENERATION_FPS):.1f}s | "
        f"num_frames={num_frames} | res={width}x{height}"
    )

    for i, prompt in enumerate(scenes_list):
        ref_label = 'last_frame' if (i > 0 and current_ref != character_image) else 'character_img'
        log_event(f"Scene {i+1}/{len(scenes_list)} | ref={ref_label} | frames={num_frames}")
        try:
            clip_bytes, last_frame_img = generate_single_scene_gpu(
                prompt=prompt,
                ref_image_path=current_ref,
                num_frames=num_frames,
                height=height,
                width=width,
                seed=seed + i,
            )

            if clip_bytes:
                clips_bytes.append(clip_bytes)

            # ── Continuity ────────────────────────────────────────────────────
            # last_frame_img is a PIL Image returned from inside the GPU function
            # (extracted BEFORE file cleanup), so it's always valid.
            if last_frame_img is not None:
                ref_path = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
                last_frame_img.save(ref_path)
                log_event(f"Scene {i+1}: last frame saved → ref for scene {i+2}")
                current_ref = ref_path
            else:
                # Keep previous ref; never reset to character_image mid-sequence
                log_event(f"Scene {i+1}: no last frame returned, keeping previous ref")

        except Exception as e:
            log_error(f"Error generating scene {i+1}: {e}")
            # current_ref unchanged → reuse last good frame for next scene

        gc.collect()

    return clips_bytes

def concatenate_videos(video_paths):
    output_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    cmd = ["ffmpeg", "-y"]
    for vp in video_paths:
        cmd.extend(["-i", vp])
    num_files = len(video_paths)
    filter_str = "".join([f"[{i}:v]fps=18,format=yuv420p,scale=768:512,setsar=1,setpts=PTS-STARTPTS[v{i}];" for i in range(num_files)])
    filter_str += "".join([f"[v{i}]" for i in range(num_files)])
    filter_str += f"concat=n={num_files}:v=1:a=0[outv]"
    cmd.extend([
        "-filter_complex", filter_str,
        "-map", "[outv]",
        "-c:v", "libx264", "-pix_fmt", "yuv420p", "-r", "18", "-vsync", "cfr", output_path
    ])
    subprocess.run(cmd, check=True)
    return output_path

def generate_movie(story_prompt, character_image, duration_per_scene, height, width, seed, randomize_seed, progress=gr.Progress()):
    curr_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
    capped_frames = compute_num_frames(duration_per_scene)
    actual_duration = (capped_frames - 1) / GENERATION_FPS
    duration_note = ""
    if actual_duration < duration_per_scene - 0.5:
        duration_note = f" (ajustado a {actual_duration:.1f}s por límite de VRAM)"

    progress(0, desc="Generando guion con Gemma-3...")
    yield [f"Generando guion con Gemma-3...{duration_note}", "", None, None, None, None, None, None, None, None, None]
    try:
        raw_script, char_desc = generate_screenplay_gemma(story_prompt, character_image)
    except Exception as e:
        yield [f"Error al generar el guion: {str(e)}", "", None, None, None, None, None, None, None, None, None]
        return
    clean_res = raw_script
    if "```" in clean_res:
        clean_res = clean_res.split("```")[1]
        if clean_res.startswith("json"):
            clean_res = clean_res[4:]
    clean_res = clean_res.strip()
    try:
        scenes = json.loads(clean_res)
    except Exception as e:
        yield ["Error al decodificar JSON del guion. Intentando recuperar...", "", None, None, None, None, None, None, None, None, None]
        scenes = []
        import re
        matches = re.findall(r'\{\s*"scene"\s*:\s*\d+\s*,\s*"prompt"\s*:\s*"[^"]*"\s*\}', clean_res)
        for m in matches:
            try:
                scenes.append(json.loads(m))
            except:
                pass
        if not scenes:
            yield ["No se pudo recuperar el guion estructurado.", "", None, None, None, None, None, None, None, None, None]
            return
    script_display = ""
    if char_desc:
        script_display += f"Personaje detectado:\n{char_desc}\n\n"
    script_display += "Guion cinematográfico:\n"
    for s in scenes:
        script_display += f"Escena {s['scene']}: {s['prompt']}\n\n"
    progress(0.1, desc=f"Guion generado. Renderizando {len(scenes[:8])} escenas ({actual_duration:.1f}s c/u)...")
    yield [f"Guion generado. Generando escenas en GPU (cada escena en su propia sesión GPU)...{duration_note}", script_display, None, None, None, None, None, None, None, None, None]
    scenes_list = [s["prompt"] for s in scenes[:8]]
    try:
        clips_bytes_list = generate_all_scenes(scenes_list, character_image, duration_per_scene, height, width, curr_seed)
    except Exception as e:
        yield [f"Error al generar escenas en GPU: {str(e)}", script_display, None, None, None, None, None, None, None, None, None]
        return
    video_clips = []
    ui_videos = [None] * 8
    for i, clip_bytes in enumerate(clips_bytes_list):
        clip_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
        with open(clip_path, "wb") as f:
            f.write(clip_bytes)
        video_clips.append(clip_path)
        ui_videos[i] = clip_path
    if not video_clips:
        yield ["Error: No se pudo generar ninguna escena.", script_display, None] + ui_videos
        return
    progress(0.9, desc="Concatenando todas las escenas en la película final...")
    yield ["Concatenando todas las escenas en la película final...", script_display, None] + ui_videos
    try:
        final_movie = concatenate_videos(video_clips)
        progress(1.0, desc="¡Película generada con éxito!")
        yield [f"¡Película generada con éxito! ({len(video_clips)} escenas × {actual_duration:.1f}s c/u)", script_display, final_movie] + ui_videos
    except Exception as e:
        yield [f"Error al concatenar videos: {str(e)}", script_display, None] + ui_videos

with gr.Blocks(title="CinemAI 🎬🎥") as demo:
    gr.Markdown("# CinemAI 🎬🎥")
    with gr.Row():
        with gr.Column():
            story_prompt = gr.Textbox(label="Idea de la Película / Historia", value="Un cuento de terror sobre una casa embrujada donde un joven explorador con chaqueta roja busca respuestas.", lines=4)
            img_in = gr.Image(label="Imagen de Referencia del Personaje", type="filepath", height=256)
            with gr.Row():
                dur_ui = gr.Slider(5.0, 30.0, 10.0, step=1.0, label="Duración por Escena (Segundos)")
                res_ui = gr.Radio(["16:9", "1:1", "9:16"], value="16:9", label="Relación de Aspecto")
            with gr.Accordion("Configuración Avanzada", open=False):
                seed = gr.Slider(0, MAX_SEED, DEFAULT_SEED, step=1, label="Semilla")
                random_seed = gr.Checkbox(label="Semilla Aleatoria", value=True)
            gen_btn = gr.Button("🎬 Generar Película Completa", variant="primary")
        with gr.Column():
            status_box = gr.Textbox(label="Estado de la Generación", interactive=False)
            video_out = gr.Video(label="Película Final Generada (Preview)", autoplay=True, height=512)
            script_out = gr.Textbox(label="Guion Cinematográfico de Gemma-3", lines=15, interactive=False)
    
    with gr.Accordion("📥 Descargar Escenas Individuales (Respaldo)", open=False):
        gr.Markdown("Si la concatenación del video final presenta problemas de duración en tu navegador, aquí puedes previsualizar y descargar cada escena por separado.")
        scene_videos = []
        with gr.Row():
            for idx in range(8):
                if idx == 4:
                    gr.HTML("<div style='flex-basis: 100%; height: 0;'></div>")
                with gr.Column(min_width=200):
                    scene_videos.append(gr.Video(label=f"Escena {idx+1}", interactive=False, height=180))
                    
    w_s, h_s = gr.State(768), gr.State(512)
    res_ui.change(lambda r: (768, 512) if r=="16:9" else ((512, 512) if r=="1:1" else (512, 768)), res_ui, [w_s, h_s])
    gen_btn.click(
        generate_movie,
        inputs=[story_prompt, img_in, dur_ui, h_s, w_s, seed, random_seed],
        outputs=[status_box, script_out, video_out] + scene_videos
    )

if __name__ == "__main__":
    demo.launch(ssr_mode=False, mcp_server=True)