#!/usr/bin/env python3
"""
Fix missing images in existing articles by assigning appropriate violation images
"""

import weaviate
import random

def select_violation_image(violations_text):
    """Select appropriate image based on violation content"""
    
    if not violations_text:
        return f'/assets/images/violations/general_{random.randint(1,3)}.jpg'
    
    violations_lower = violations_text.lower()
    
    # Map violation keywords to specific images
    if any(word in violations_lower for word in ['temperature', 'temp', 'cold', 'hot', 'tcs']):
        return f'/assets/images/violations/temperature_{random.randint(1,3)}.jpg'
    elif any(word in violations_lower for word in ['rodent', 'mouse', 'mice', 'rat', 'droppings']):
        return f'/assets/images/violations/rodent_{random.randint(1,3)}.jpg'
    elif any(word in violations_lower for word in ['roach', 'cockroach', 'insect', 'fly', 'flies']):
        return f'/assets/images/violations/pest_{random.randint(1,3)}.jpg'
    elif any(word in violations_lower for word in ['hand', 'wash', 'soap', 'sanitizer', 'hygiene']):
        return f'/assets/images/violations/handwashing_{random.randint(1,3)}.jpg'
    elif any(word in violations_lower for word in ['clean', 'sanitize', 'dirty', 'soil', 'grease']):
        return f'/assets/images/violations/sanitation_{random.randint(1,3)}.jpg'
    elif any(word in violations_lower for word in ['storage', 'store', 'label', 'date', 'fifo']):
        return f'/assets/images/violations/storage_{random.randint(1,3)}.jpg'
    elif any(word in violations_lower for word in ['cross', 'contamination', 'raw', 'cooked']):
        return f'/assets/images/violations/cross_contamination_{random.randint(1,3)}.jpg'
    elif any(word in violations_lower for word in ['license', 'permit', 'certificate', 'document']):
        return f'/assets/images/violations/documentation_{random.randint(1,3)}.jpg'
    elif any(word in violations_lower for word in ['floor', 'wall', 'ceiling', 'repair', 'maintain']):
        return f'/assets/images/violations/structural_{random.randint(1,3)}.jpg'
    elif any(word in violations_lower for word in ['plumb', 'water', 'sink', 'drain', 'leak']):
        return f'/assets/images/violations/plumbing_{random.randint(1,3)}.jpg'
    elif 'closure' in violations_lower or 'closed' in violations_lower:
        return f'/assets/images/violations/closure_{random.randint(1,3)}.jpg'
    else:
        # Default to general violation image
        return f'/assets/images/violations/general_{random.randint(1,3)}.jpg'

def main():
    print("🔧 Fixing article images...")
    
    # Connect to Weaviate
    client = weaviate.connect_to_local()
    
    try:
        collection = client.collections.get('Articles')
        
        # Get all articles
        print("📊 Fetching all articles...")
        articles = collection.query.fetch_objects(
            limit=1000,
            return_properties=['title', 'violations', 'image_url']
        )
        
        print(f"Found {len(articles.objects)} articles to process")
        
        updated_count = 0
        skipped_count = 0
        
        for article in articles.objects:
            article_id = article.uuid
            props = article.properties
            
            # Skip if already has an image
            if props.get('image_url') and props['image_url'] != 'None':
                skipped_count += 1
                continue
            
            # Select appropriate image based on violations
            violations = props.get('violations', '')
            image_url = select_violation_image(violations)
            
            # Update the article
            try:
                collection.data.update(
                    uuid=article_id,
                    properties={'image_url': image_url}
                )
                updated_count += 1
                print(f"✅ Updated: {props['title'][:50]}...")
                print(f"   Image: {image_url}")
                
            except Exception as e:
                print(f"❌ Failed to update article: {e}")
        
        print(f"\n✨ Complete!")
        print(f"   Updated: {updated_count} articles")
        print(f"   Skipped: {skipped_count} articles (already had images)")
        
    finally:
        client.close()

if __name__ == "__main__":
    main()