import json
import os
import torch
from diffusers import AutoPipelineForText2Image

# Configuration
MANIFEST_PATH = '/Users/apayne/the-testament/visual_manifest.json'
OUTPUT_DIR = '/Users/apayne/the-testament/assets/visuals'
MODEL_ID = "stabilityai/sd-turbo"

def load_manifest():
    with open(MANIFEST_PATH, 'r') as f:
        return json.load(f)

def save_manifest(manifest):
    with open(MANIFEST_PATH, 'w') as f:
        json.dump(manifest, f, indent=2)

def run_batch_render(limit=10):
    manifest = load_manifest()
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    
    print(f"Loading {MODEL_ID} on MPS...")
    pipe = AutoPipelineForText2Image.from_pretrained(
        MODEL_ID, torch_dtype=torch.float16, variant="fp16"
    ).to("mps")

    count = 0
    for beat in manifest:
        if count >= limit:
            break
        if beat.get('status') == 'completed':
            continue

        text = beat['text']
        prompt = f"Cinematic film still, moody noir, {text[:150]}, 35mm, muted tones, high contrast, detailed textures"
        
        filename = f"{beat['chapter'].replace(' ', '_').replace('#', '').replace('—', '_')}_p{beat['paragraph']}.png"
        path = os.path.join(OUTPUT_DIR, filename)

        try:
            image = pipe(prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0]
            image.save(path)
            beat['status'] = 'completed'
            beat['final_prompt'] = prompt
            beat['path'] = path
            count += 1
        except Exception as e:
            print(f"Error rendering {filename}: {e}")

    save_manifest(manifest)
    print(f"Batch complete. Rendered {count} images.")

if __name__ == "__main__":
    run_batch_render(limit=5)