#!/usr/bin/env python3
"""
Reprocess existing images to add gradient metadata.
This script reads all existing images in the output folder and generates
gradient backgrounds for them, appending to metadata.jsonl
"""

import numpy as np
import imageio.v3 as iio
import os
import re
import json
from tqdm import tqdm

OUTPUT_DIR = "output"

def img_to_gradient(input_image):
    """Generate CSS gradient string from image"""
    string = "linear-gradient(in oklab"
    for i in np.array_split(input_image[:,0,:], 11):
        r, g, b = np.mean(i, axis=0).astype(np.uint8)
        string += f",rgb({r},{g},{b})"
    string += ")"
    return string

def extract_length(filename):
    """Extract length from filename like 'live-timestamp-start-end-length.ext'"""
    match = re.match(r'live-\d+-\d+-\d+-(\d+)\.', filename)
    if match:
        return int(match.group(1))
    return None

def main():
    # Get all existing image files
    image_files = []
    if os.path.exists(OUTPUT_DIR):
        image_files = [f for f in os.listdir(OUTPUT_DIR) 
                      if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp', '.gif', '.avif'))]
    
    if not image_files:
        print("No images found in output directory")
        return
    
    print(f"Found {len(image_files)} images to process")
    
    # Load existing metadata to avoid duplicates
    existing_metadata = set()
    metadata_path = os.path.join(OUTPUT_DIR, 'metadata.jsonl')
    if os.path.exists(metadata_path):
        with open(metadata_path, 'r') as f:
            for line in f:
                if line.strip():
                    data = json.loads(line)
                    existing_metadata.add(data['filename'])
        print(f"Found {len(existing_metadata)} existing metadata entries")
    
    # Process each image
    processed = 0
    skipped = 0
    
    for filename in tqdm(image_files, desc="Processing images"):
        # Skip if metadata already exists
        if filename in existing_metadata:
            skipped += 1
            continue
        
        try:
            # Read the image
            filepath = os.path.join(OUTPUT_DIR, filename)
            image = iio.imread(filepath)
            
            # Rotate back to vertical orientation (undo the rot90 from write)
            image_vertical = image
            
            # Generate gradient
            gradient_string = img_to_gradient(image_vertical)
            
            # Extract length from filename
            length = extract_length(filename)
            if length is None:
                print(f"Warning: Could not extract length from {filename}")
                length = image.shape[1]  # Use actual width as fallback
            
            # Append to metadata file
            with open(metadata_path, 'a') as f:
                metadata_entry = {
                    "filename": filename,
                    "length": length,
                    "gradient": gradient_string
                }
                f.write(json.dumps(metadata_entry) + '\n')
            
            processed += 1
            
        except Exception as e:
            print(f"Error processing {filename}: {e}")
    
    print(f"\nProcessing complete!")
    print(f"Processed: {processed}")
    print(f"Skipped (already had metadata): {skipped}")
    print(f"Total metadata entries: {processed + len(existing_metadata)}")

if __name__ == '__main__':
    main()