"""
Script to iterate over all images in the img folder and insert their metadata into the avdb_image table.
"""
import os
from PIL import Image
import sqlalchemy
from sqlalchemy import text, exc

# Database URL from alembic.ini
DATABASE_URL = "postgresql+psycopg2://postgres:pgRocks@localhost:5432/testdb"

def get_image_metadata(image_path):
    """Extract metadata (width, height) from an image."""
    with Image.open(image_path) as img:
        width, height = img.size
    return width, height

def insert_images_into_db():
    """Iterate over all images in the img folder and insert their metadata into the avdb_image table."""
    engine = sqlalchemy.create_engine(DATABASE_URL)

    # Iterate over all subdirectories in the img folder
    img_folder = "img"
    inserted_count = 0
    skipped_count = 0
    test_shim = 'TESTIMG.jpeg'
    duplicate_count = 0
    
    for root, dirs, files in os.walk(img_folder):
        for file in files:
            # test nmode
            if file.endswith(test_shim):
            # if file.lower().endswith(('.jpg', '.jpeg')):
                image_path = os.path.join(root, file)
                img_name = file
                
                # Extract width and height
                try:
                    width, height = get_image_metadata(image_path)
                except Exception as e:
                    print(f"Failed to extract metadata for {image_path}: {e}")
                    skipped_count += 1
                    continue
                
                # Insert into the database
                try:
                    with engine.connect() as connection:
                        query = text(
                            """
                            INSERT INTO avdb_image (img_path, img_name, width, height)
                            VALUES (:img_path, :img_name, :width, :height)
                            """
                        )
                        connection.execute(
                            query,
                            {"img_path": image_path, "img_name": img_name, "width": width, "height": height}
                        )
                        connection.commit()
                        inserted_count += 1
                        print(f"Inserted: {img_name} (Width: {width}, Height: {height})")
                except exc.IntegrityError:
                    # Handle duplicate key violations
                    print(f"Duplicate entry skipped for {image_path}")
                    duplicate_count += 1
                except Exception as e:
                    print(f"Failed to insert {image_path} into the database: {e}")
                    skipped_count += 1
    
    print(f"\nTotal images inserted: {inserted_count}")
    print(f"Total duplicates skipped: {duplicate_count}")
    print(f"Total images skipped due to errors: {skipped_count}")

if __name__ == "__main__":
    insert_images_into_db()