import config, os
from PIL import Image
import imagehash

def generate_phash(image_path):
    image = Image.open(image_path)
    return str(imagehash.phash(image))

db, cursor = config.gen_connection()

cursor.execute("SELECT id, media_id, media_url FROM media WHERE media_type = 'image' AND phash IS NULL;")
results = cursor.fetchall()

count = 0
cacheDir = 'cache'
os.makedirs(cacheDir, exist_ok=True)
print(f"Found {len(results)} files to process.")


for result in results:
    count += 1
    itemID = result[0]
    mediaID = result[1]
    mediaURL = result[2]
    
    serverPath = mediaURL.replace("https://storysave.b-cdn.net/", '').replace('//', '/').replace('\\', '/')
    localFilePath = os.path.join(cacheDir, os.path.basename(serverPath))
    
    if not os.path.exists(localFilePath):
        print(f"File {localFilePath} does not exist, skipping.")
        continue
    
    try:
        phash = generate_phash(localFilePath)
        
        cursor.execute("UPDATE media SET phash = %s WHERE id = %s", (phash, itemID))
        db.commit()
        
        print(f"Processed {count}/{len(results)}: {mediaID} with pHash {phash}")
    except Exception as e:
        print(f"Error processing {mediaID}: {e}")