import json
import random
import os
from diffusers import StableDiffusionPipeline
import torch
import openai

# Load and cache the diffusion pipeline (only once)
pipe = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    torch_dtype=torch.float16
)
pipe = pipe.to("cpu")


openai.api_key = os.getenv("OPENAI_API_KEY")  # Make sure this is set in your environment

# Global story context (in Chinese)
story_context_cn = "《博物馆的全能ACE》是一部拟人化博物馆文物与AI讲解助手互动的短片，讲述太阳人石刻在闭馆后的博物馆中，遇到了新来的AI助手博小翼，两者展开对话，AI展示了自己的多模态讲解能力与文化知识，最终被文物们认可，并一起展开智慧导览服务的故事。该片融合了文物拟人化、夜间博物馆奇妙氛围、科技感界面与中国地方文化元素，风格活泼、具未来感。"

def generate_keyframe_prompt(segment):
    """
    Calls GPT-4o to generate an image prompt optimized for Stable Diffusion,
    based on segment content and full story context.
    """
    description = segment.get("description", "")
    speaker = segment.get("speaker", "")
    narration = segment.get("narration", "")
    segment_id = segment.get("segment_id")

    input_prompt = f"你是一个擅长视觉脚本设计的AI，请基于以下故事整体背景与分镜内容，帮我生成一个适合用于Stable Diffusion图像生成的英文提示词（image prompt），用于生成低分辨率草图风格的关键帧。请注意突出主要角色、镜头氛围、光影、构图、动作，避免复杂背景和细节。

【整体故事背景】：\n{story_context_cn}

【当前分镜描述】：\n{description}
【角色】：{speaker}\n【台词或画外音】：{narration}

请用英文输出一个简洁但具体的prompt，风格偏草图、线稿、卡通、简洁构图，并指出一个negative prompt。"

    try:
        response = openai.ChatCompletion.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": "You are an expert visual prompt designer for image generation."},
                {"role": "user", "content": input_prompt}
            ],
            temperature=0.7
        )
        output_text = response["choices"][0]["message"]["content"]

        # Split response into prompt + negative if possible
        if "Negative prompt:" in output_text:
            prompt, negative = output_text.split("Negative prompt:", 1)
        else:
            prompt, negative = output_text, "blurry, distorted, low quality, text, watermark"

        return {
            "prompt": prompt.strip(),
            "negative_prompt": negative.strip()
        }
    except Exception as e:
        print(f"[Error] GPT-4o prompt generation failed for segment {segment_id}: {e}")
        return {
            "prompt": description,
            "negative_prompt": ""
        }


def generate_all_keyframe_images(script_data, output_dir="keyframes"):
    """
    Generates 3 keyframe images per segment using Stable Diffusion,
    stores them in the given output directory.
    """
    os.makedirs(output_dir, exist_ok=True)
    keyframe_outputs = []

    for segment in script_data:
        sd_prompts = generate_keyframe_prompt(segment)
        prompt = sd_prompts["prompt"]
        negative_prompt = sd_prompts["negative_prompt"]
        segment_id = segment.get("segment_id")

        frame_images = []
        for i in range(3):
            image = pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=20, guidance_scale=7.5, height=256, width=256).images[0]
            image_path = os.path.join(output_dir, f"segment_{segment_id}_v{i+1}.png")
            image.save(image_path)
            frame_images.append(image_path)

        keyframe_outputs.append({
            "segment_id": segment_id,
            "prompt": prompt,
            "negative_prompt": negative_prompt,
            "frame_images": frame_images
        })

        print(f"✓ Generated 3 images for Segment {segment_id}")

    return keyframe_outputs