"""
🎬 FULL AI PIPELINE HORROR SHORTS GENERATOR
Everything AI-Generated: Story → Speech → Images → Video

PIPELINE:
1. 🤖 LLM writes horror story (Mistral-7B)
2. 🎙️ AI generates speech (Bark TTS)
3. 🎨 AI creates images (Stable Diffusion XL)
4. 🎵 AI generates ambient sound
5. 🎬 Combines into final video

100% Free Hugging Face Models - No API Keys Needed
"""

import gradio as gr
import torch
import random
import numpy as np
import cv2
from PIL import Image, ImageDraw, ImageFont, ImageEnhance
import os
import shutil
import gc
import re
from typing import List, Tuple

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
from diffusers import StableDiffusionXLPipeline, DPMSolverMultistepScheduler
from bark import SAMPLE_RATE, generate_audio, preload_models
from scipy.io.wavfile import write as write_wav
from pydub import AudioSegment
from pydub.generators import Sine, WhiteNoise

# ═══════════════════════════════════════════════════════════════════
# STEP 1: AI STORY GENERATION (LLM)
# ═══════════════════════════════════════════════════════════════════

_llm_model = None
_llm_tokenizer = None

def load_story_llm():
    """Load Mistral-7B for story generation."""
    global _llm_model, _llm_tokenizer
    
    if _llm_model is None:
        print("Loading Mistral-7B for story generation...")
        
        model_name = "mistralai/Mistral-7B-Instruct-v0.2"
        
        _llm_tokenizer = AutoTokenizer.from_pretrained(model_name)
        _llm_model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto" if torch.cuda.is_available() else None,
            low_cpu_mem_usage=True
        )
        
        print("Story LLM loaded!")
    
    return _llm_model, _llm_tokenizer

def generate_horror_story_with_ai(theme: str = None) -> dict:
    """Use LLM to generate original horror story."""
    
    model, tokenizer = load_story_llm()
    
    # Themes for variety
    themes = [
        "liminal spaces and parallel dimensions",
        "time loops and paradoxes", 
        "surveillance and being watched",
        "mirrors and reflections",
        "abandoned buildings with secrets",
        "technology that behaves impossibly"
    ]
    
    if theme is None:
        theme = random.choice(themes)
    
    # Prompt engineered for horror stories with loops
    prompt = f"""[INST] You are a master horror writer specializing in creepypasta and internet horror.

Write a SHORT horror story (exactly 250-300 words) with these requirements:

THEME: {theme}
STYLE: First-person narration, present tense, internet creepypasta
STRUCTURE: 
- Hook in first sentence
- Build tension gradually
- End with a twist that CONNECTS BACK to the beginning (looping narrative)
- The ending should make the reader want to re-read from the start

TONE: Unsettling, atmospheric, psychological horror (not gore)
AVOID: Clichés, explaining too much, happy endings

Write the story now (250-300 words): [/INST]

"""
    
    inputs = tokenizer(prompt, return_tensors="pt")
    if torch.cuda.is_available():
        inputs = inputs.to("cuda")
    
    outputs = model.generate(
        **inputs,
        max_new_tokens=400,
        temperature=0.8,
        top_p=0.9,
        do_sample=True,
        repetition_penalty=1.15
    )
    
    story = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Extract just the story (remove prompt)
    story = story.split("[/INST]")[-1].strip()
    
    # Clean up
    story = re.sub(r'\n\n+', '\n\n', story)
    
    # Generate title with AI
    title_prompt = f"[INST] Give a 2-4 word creepy title for this horror story: {story[:100]}... [/INST] Title:"
    
    title_inputs = tokenizer(title_prompt, return_tensors="pt")
    if torch.cuda.is_available():
        title_inputs = title_inputs.to("cuda")
    
    title_outputs = model.generate(
        **title_inputs,
        max_new_tokens=10,
        temperature=0.7
    )
    
    title = tokenizer.decode(title_outputs[0], skip_special_tokens=True)
    title = title.split("Title:")[-1].strip().split("\n")[0]
    title = re.sub(r'[^a-zA-Z0-9\s]', '', title)[:50]
    
    # Generate scene descriptions
    scene_prompts = generate_scene_descriptions_from_story(story)
    
    return {
        "title": title if title else "Untitled Horror",
        "script": story,
        "theme": theme,
        "scene_prompts": scene_prompts
    }

def generate_scene_descriptions_from_story(story: str) -> List[str]:
    """Extract key moments and generate visual prompts."""
    
    # Split story into roughly 8-10 segments
    sentences = [s.strip() for s in re.split(r'[.!?]+', story) if s.strip()]
    
    # Group into scenes
    scenes_per_segment = max(1, len(sentences) // 8)
    scene_groups = [sentences[i:i+scenes_per_segment] for i in range(0, len(sentences), scenes_per_segment)]
    
    # Generate visual prompts based on content
    prompts = []
    
    for group in scene_groups[:10]:  # Max 10 scenes
        text = ' '.join(group).lower()
        
        # Keyword-based scene generation
        if any(word in text for word in ['door', 'entrance', 'hallway']):
            prompts.append("mysterious door in dark hallway, ominous atmosphere, cinematic lighting, horror aesthetic")
        elif any(word in text for word in ['mirror', 'reflection', 'glass']):
            prompts.append("eerie mirror reflection, bathroom, dim lighting, unsettling atmosphere, horror movie")
        elif any(word in text for word in ['stair', 'stairs', 'staircase']):
            prompts.append("dark staircase, shadows, ominous perspective, horror atmosphere, dramatic lighting")
        elif any(word in text for word in ['window', 'outside', 'view']):
            prompts.append("view through window, ominous sky, dramatic lighting, horror atmosphere, cinematic")
        elif any(word in text for word in ['room', 'apartment', 'house']):
            prompts.append("empty room, liminal space, eerie atmosphere, dramatic shadows, horror aesthetic")
        elif any(word in text for word in ['forest', 'woods', 'trees']):
            prompts.append("dark forest, fog, mysterious atmosphere, horror movie lighting, cinematic")
        elif any(word in text for word in ['camera', 'footage', 'monitor']):
            prompts.append("security camera footage, grainy, CCTV aesthetic, surveillance horror, dramatic")
        elif any(word in text for word in ['elevator', 'floor']):
            prompts.append("elevator interior, flickering lights, claustrophobic, horror atmosphere, cinematic")
        else:
            prompts.append("dark atmospheric horror scene, liminal space, eerie lighting, unsettling, cinematic")
    
    # Ensure we have at least 8 prompts
    while len(prompts) < 8:
        prompts.append("abstract horror atmosphere, darkness, shadows, eerie mood, cinematic lighting")
    
    return prompts[:10]

# ═══════════════════════════════════════════════════════════════════
# STEP 2: AI SPEECH GENERATION (BARK TTS)
# ═══════════════════════════════════════════════════════════════════

def load_bark_tts():
    """Load Bark TTS model."""
    print("Loading Bark TTS...")
    preload_models()
    print("Bark TTS ready!")

def generate_ai_speech(text: str, target_duration: float = 55.0) -> Tuple[str, float]:
    """Generate speech with Bark AI TTS."""
    
    load_bark_tts()
    
    # Bark works best with shorter segments
    # Split text into chunks
    sentences = [s.strip() + '.' for s in re.split(r'[.!?]+', text) if s.strip()]
    
    audio_segments = []
    
    print(f"Generating speech for {len(sentences)} sentences...")
    
    for i, sentence in enumerate(sentences):
        print(f"  Generating sentence {i+1}/{len(sentences)}...")
        
        # Generate audio with Bark
        # Use a creepy voice preset
        audio_array = generate_audio(
            sentence,
            history_prompt="v2/en_speaker_6",  # Deeper, more ominous voice
        )
        
        # Convert to AudioSegment
        temp_path = f"temp/bark_segment_{i}.wav"
        write_wav(temp_path, SAMPLE_RATE, audio_array)
        
        segment = AudioSegment.from_wav(temp_path)
        audio_segments.append(segment)
        
        # Cleanup
        os.remove(temp_path)
    
    # Combine all segments
    full_audio = sum(audio_segments)
    
    # Adjust speed to hit target duration
    current_duration = len(full_audio) / 1000.0
    
    if abs(current_duration - target_duration) > 2:
        speed_factor = current_duration / target_duration
        full_audio = full_audio._spawn(
            full_audio.raw_data,
            overrides={"frame_rate": int(full_audio.frame_rate * speed_factor)}
        ).set_frame_rate(SAMPLE_RATE)
    
    # Horror audio processing
    full_audio = full_audio - 2  # Slight reduction
    
    # Add reverb
    reverb = full_audio - 20
    full_audio = full_audio.overlay(reverb, position=70)
    
    # Fades
    full_audio = full_audio.fade_in(300).fade_out(500)
    
    # Force to exactly target duration
    full_audio = full_audio[:int(target_duration * 1000)]
    
    # Export
    output_path = "temp/ai_voice.mp3"
    full_audio.export(output_path, format='mp3', bitrate="192k")
    
    return output_path, len(full_audio) / 1000.0

# ═══════════════════════════════════════════════════════════════════
# STEP 3: AI IMAGE GENERATION (SDXL)
# ═══════════════════════════════════════════════════════════════════

_sdxl_pipe = None

def load_image_generator():
    """Load SDXL for image generation."""
    global _sdxl_pipe
    
    if _sdxl_pipe is None:
        print("Loading Stable Diffusion XL...")
        
        _sdxl_pipe = StableDiffusionXLPipeline.from_pretrained(
            "stabilityai/stable-diffusion-xl-base-1.0",
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            use_safetensors=True,
            variant="fp16" if torch.cuda.is_available() else None
        )
        
        _sdxl_pipe.scheduler = DPMSolverMultistepScheduler.from_config(
            _sdxl_pipe.scheduler.config
        )
        
        if torch.cuda.is_available():
            _sdxl_pipe.to("cuda")
            _sdxl_pipe.enable_vae_slicing()
        else:
            _sdxl_pipe.enable_attention_slicing()
            _sdxl_pipe.enable_vae_slicing()
        
        print("SDXL ready!")
    
    return _sdxl_pipe

def generate_ai_image(prompt: str, index: int) -> Image.Image:
    """Generate image with AI."""
    
    pipe = load_image_generator()
    
    image = pipe(
        prompt=prompt + ", cinematic, dramatic lighting, horror atmosphere, high quality, professional",
        negative_prompt="blurry, low quality, text, watermark, bright, cheerful, cartoon",
        num_inference_steps=25,
        guidance_scale=7.5,
        height=1024,
        width=768,
    ).images[0]
    
    # Apply horror grading
    enhancer = ImageEnhance.Color(image)
    image = enhancer.enhance(0.4)
    
    enhancer = ImageEnhance.Contrast(image)
    image = enhancer.enhance(1.4)
    
    enhancer = ImageEnhance.Brightness(image)
    image = enhancer.enhance(0.75)
    
    # Clear memory
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    gc.collect()
    
    return image

# ═══════════════════════════════════════════════════════════════════
# STEP 4: VIDEO ASSEMBLY
# ═══════════════════════════════════════════════════════════════════

def setup_dirs():
    for folder in ['output', 'temp', 'images']:
        if os.path.exists(folder):
            shutil.rmtree(folder)
        os.makedirs(folder)

def create_ambient_sound(duration: float) -> str:
    """Generate AI-like ambient sound."""
    duration_ms = int(duration * 1000)
    
    # Multi-layer ambient
    drone1 = Sine(55).to_audio_segment(duration=duration_ms) - 20
    drone2 = Sine(110).to_audio_segment(duration=duration_ms) - 23
    tension = Sine(8000).to_audio_segment(duration=duration_ms) - 30
    noise = WhiteNoise().to_audio_segment(duration=duration_ms) - 35
    
    ambient = drone1.overlay(drone2).overlay(tension).overlay(noise)
    ambient = ambient.fade_in(3000).fade_out(3000)
    
    ambient.export("temp/ambient.mp3", format='mp3')
    return "temp/ambient.mp3"

def animate_image(img: Image.Image, duration: float, movement: str) -> List[np.ndarray]:
    """Create animation from image."""
    arr = np.array(img)
    arr = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
    
    h, w = arr.shape[:2]
    frames = []
    total_frames = int(duration * 30)
    
    # Scale for movement
    scaled = cv2.resize(arr, (int(w*1.3), int(h*1.3)), interpolation=cv2.INTER_LINEAR)
    sh, sw = scaled.shape[:2]
    
    for i in range(total_frames):
        progress = i / total_frames
        ease = progress * progress * (3.0 - 2.0 * progress)
        
        if movement == 'zoom':
            s = 1.0 + ease * 0.2
            temp = cv2.resize(arr, (int(w*s), int(h*s)), interpolation=cv2.INTER_LINEAR)
            th, tw = temp.shape[:2]
            x, y = (tw-w)//2, (th-h)//2
            frame = temp[y:y+h, x:x+w]
        else:  # pan
            x = int((sw-w) * ease)
            frame = scaled[0:h, x:x+w]
        
        frames.append(frame)
    
    return frames

def upscale_frame(frame: np.ndarray) -> np.ndarray:
    """Upscale to 1080x1920."""
    target_w, target_h = 1080, 1920
    h, w = frame.shape[:2]
    
    scale = max(target_w/w, target_h/h)
    new_size = (int(w*scale), int(h*scale))
    
    upscaled = cv2.resize(frame, new_size, interpolation=cv2.INTER_LANCZOS4)
    
    uh, uw = upscaled.shape[:2]
    x = (uw - target_w) // 2
    y = (uh - target_h) // 2
    
    return upscaled[y:y+target_h, x:x+target_w]

def add_subtitles(frame: np.ndarray, text: str) -> np.ndarray:
    """Add subtitles to frame."""
    rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    pil_img = Image.fromarray(rgb)
    draw = ImageDraw.Draw(pil_img)
    
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 55)
    except:
        font = ImageFont.load_default()
    
    # Word wrap
    words = text.split()
    lines = []
    current = []
    
    for word in words:
        test = ' '.join(current + [word])
        bbox = draw.textbbox((0, 0), test, font=font)
        if bbox[2] - bbox[0] <= 980:
            current.append(word)
        else:
            if current:
                lines.append(' '.join(current))
            current = [word]
    if current:
        lines.append(' '.join(current))
    
    # Draw
    y = 1700
    for line in lines[:2]:  # Max 2 lines
        bbox = draw.textbbox((0, 0), line, font=font)
        x = (1080 - (bbox[2] - bbox[0])) // 2
        
        # Outline
        for dx in [-4, 0, 4]:
            for dy in [-4, 0, 4]:
                draw.text((x+dx, y+dy), line, font=font, fill='black')
        
        draw.text((x, y), line, font=font, fill='white')
        y += 70
    
    return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)

def render_video(frames: List[np.ndarray], voice: str, ambient: str, output: str) -> str:
    """Render final video."""
    temp_vid = "temp/video.mp4"
    
    out = cv2.VideoWriter(temp_vid, cv2.VideoWriter_fourcc(*'mp4v'), 30, (1080, 1920))
    for f in frames:
        out.write(f)
    out.release()
    
    # Mix audio
    v = AudioSegment.from_mp3(voice)
    a = AudioSegment.from_mp3(ambient)
    mixed = v.overlay(a - 15)
    mixed = mixed[:55000]  # Exactly 55s
    mixed.export("temp/audio.mp3", format='mp3')
    
    # Combine
    cmd = f'ffmpeg -y -i {temp_vid} -i temp/audio.mp3 -c:v libx264 -preset medium -crf 20 -c:a aac -b:a 192k -t 55 -shortest {output} -loglevel error'
    os.system(cmd)
    
    return output

# ═══════════════════════════════════════════════════════════════════
# MAIN PIPELINE
# ═══════════════════════════════════════════════════════════════════

def generate_full_ai_pipeline(selected_theme: str = "Random", progress=gr.Progress()):
    """
    Complete AI pipeline: Story → Speech → Images → Video
    """
    
    try:
        setup_dirs()
        
        # STEP 1: AI writes story
        progress(0.05, desc="🤖 AI writing horror story...")
        
        theme = None if selected_theme == "Random" else selected_theme
        story_data = generate_horror_story_with_ai(theme)
        
        title = story_data['title']
        script = story_data['script']
        scene_prompts = story_data['scene_prompts']
        
        progress(0.15, desc=f"✅ Story complete: '{title}'")
        
        # STEP 2: AI generates speech
        progress(0.20, desc="🎙️ AI generating speech with Bark...")
        voice_path, duration = generate_ai_speech(script, 55.0)
        
        progress(0.35, desc=f"✅ Speech generated ({duration:.1f}s)")
        
        # STEP 3: Generate ambient
        progress(0.40, desc="🎵 Creating ambient soundscape...")
        ambient_path = create_ambient_sound(55.0)
        
        # STEP 4: AI generates images
        progress(0.45, desc="🎨 Loading image AI...")
        load_image_generator()
        
        num_scenes = min(len(scene_prompts), 8)
        sec_per_scene = 55.0 / num_scenes
        all_frames = []
        
        movements = ['zoom', 'pan'] * 5
        
        for i in range(num_scenes):
            progress(0.45 + (i * 0.05), desc=f"🎨 AI generating image {i+1}/{num_scenes}...")
            
            img = generate_ai_image(scene_prompts[i], i)
            
            progress(0.45 + (i * 0.05) + 0.02, desc=f"🎞️ Animating scene {i+1}/{num_scenes}...")
            
            frames = animate_image(img, sec_per_scene, movements[i])
            frames = [upscale_frame(f) for f in frames]
            
            all_frames.extend(frames)
            
            del img, frames
            gc.collect()
        
        # STEP 5: Add subtitles
        progress(0.90, desc="📄 Adding subtitles...")
        
        sentences = [s.strip() + '.' for s in re.split(r'[.!?]+', script) if s.strip()]
        frames_per_sub = len(all_frames) // len(sentences)
        
        final_frames = []
        for i, frame in enumerate(all_frames):
            sub_idx = min(i // frames_per_sub, len(sentences) - 1)
            final_frames.append(add_subtitles(frame, sentences[sub_idx]))
        
        # STEP 6: Render
        progress(0.95, desc="🎬 Rendering final video...")
        output = render_video(final_frames, voice_path, ambient_path, "output/ai_horror_short.mp4")
        
        progress(1.0, desc="✅ Complete!")
        
        info = f"""
### 🤖 Full AI Generation Complete!

**Title:** {title}

**AI Pipeline:**
1. ✅ Story written by: Mistral-7B-Instruct
2. ✅ Speech by: Bark TTS (Suno AI)
3. ✅ Images by: Stable Diffusion XL
4. ✅ Assembled automatically

**Stats:**
- Duration: 55.0 seconds
- Scenes: {num_scenes}
- Frames: {len(final_frames)}
- Theme: {story_data['theme']}

**Everything created by AI - zero human writing!**
"""
        
        return output, script, info
        
    except Exception as e:
        error = f"❌ Error: {str(e)}"
        print(error)
        import traceback
        traceback.print_exc()
        return None, error, error

# ═══════════════════════════════════════════════════════════════════
# GRADIO INTERFACE
# ═══════════════════════════════════════════════════════════════════

with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", secondary_hue="slate")) as demo:
    
    gr.Markdown("""
    # 🤖 Full AI Horror Shorts Pipeline
    ## Every Step Generated by AI - Story to Final Video
    
    **100% AI-Generated Content Using Free Hugging Face Models**
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            
            theme_dropdown = gr.Dropdown(
                choices=[
                    "Random",
                    "liminal spaces and parallel dimensions",
                    "time loops and paradoxes",
                    "surveillance and being watched",
                    "mirrors and reflections",
                    "abandoned buildings with secrets",
                    "technology that behaves impossibly"
                ],
                value="Random",
                label="🎭 Story Theme"
            )
            
            generate_btn = gr.Button(
                "🤖 Generate Full AI Horror Short",
                variant="primary",
                size="lg"
            )
            
            gr.Markdown("""
            ### 🔄 AI Pipeline Steps:
            
            **1. Story Generation** 🤖
            - Model: Mistral-7B-Instruct
            - Writes original 250-300 word story
            - Creates looping narrative
            - Generates title
            
            **2. Speech Synthesis** 🎙️
            - Model: Bark TTS (Suno AI)
            - Natural-sounding voice
            - Horror audio processing
            - Exactly 55 seconds
            
            **3. Image Generation** 🎨
            - Model: Stable Diffusion XL
            - 8 unique horror scenes
            - Cinematic color grading
            - High resolution
            
            **4. Video Assembly** 🎬
            - Animated camera movements
            - Professional subtitles
            - Layered ambient sound
            - 1080x1920 output
            
            ### ⏱️ Generation Time:
            - Story: 1-2 min
            - Speech: 3-5 min
            - Images: 20-30 min (8 scenes)
            - Assembly: 2-3 min
            
            **Total: 30-40 minutes**
            
            ### 💡 Features:
            - ✅ Zero pre-written content
            - ✅ Every story is unique
            - ✅ Free HuggingFace models
            - ✅ No API keys needed
            - ✅ Looping narratives
            - ✅ Professional quality
            """)
        
        with gr.Column(scale=2):
            video_output = gr.Video(
                label="🎬 AI-Generated Horror Short",
                height=750
            )
            
            script_output = gr.Textbox(
                label="📝 AI-Written Story",
                lines=15
            )
            
            info_output = gr.Markdown(label="📊 Generation Info")
    
    generate_btn.click(
        fn=generate_full_ai_pipeline,
        inputs=[theme_dropdown],
        outputs=[video_output, script_output, info_output]
    )
    
    gr.Markdown("""
    ---
    
    ## 🚀 Models Used (All Free from Hugging Face):
    
    1. **Mistral-7B-Instruct-v0.2** - Story generation
       - 7 billion parameters
       - Instruction-tuned for creative writing
       - Excellent at horror narratives
    
    2. **Bark TTS** - Speech synthesis
       - By Suno AI
       - Natural prosody and emotion
       - Multiple voice options
    
    3. **Stable Diffusion XL** - Image generation
       - State-of-the-art image quality
       - 1024px native resolution
       - Excellent at atmospheric scenes
    
    ## 📦 Requirements:
    
    ```
    gradio
    torch
    transformers
    diffusers
    accelerate
    bark
    scipy
    pydub
    opencv-python-headless
    pillow
    numpy
    ```
    
    ## 🎯 Best Practices:
    
    - Use GPU for reasonable speed (30-40 min)
    - CPU will work but take 2-3 hours
    - First run downloads models (~15GB total)
    - Subsequent runs use cached models
    
    ## 💰 Cost:
    
    **$0** - Completely free!
    - All models from Hugging Face
    - No API keys or subscriptions
    - Run on free GPU (Google Colab, HF Spaces)
    
    ## 🎨 Why This Is Special:
    
    Most "AI video generators" use:
    - Pre-written scripts ❌
    - Pre-recorded voice ❌
    - Stock images ❌
    
    This uses:
    - AI-written stories ✅
    - AI-generated speech ✅
    - AI-generated images ✅
    
    **Every single element created by AI!**
    """)

if __name__ == "__main__":
    demo.launch()

"""
═══════════════════════════════════════════════════════════════════
🤖 FULL AI PIPELINE - NO HUMAN INPUT REQUIRED
═══════════════════════════════════════════════════════════════════

This is a TRUE end-to-end AI content generation pipeline.

STEP 1: LLM writes story (Mistral-7B)
STEP 2: TTS creates speech (Bark)
STEP 3: Diffusion creates images (SDXL)
STEP 4: Assembly creates video

Everything automated. Every video unique. Zero templates.

Deploy on HuggingFace Spaces with GPU for best results!
═══════════════════════════════════════════════════════════════════
"""