"""
Extract weights from .keras file without deserializing the model architecture.
.keras files are ZIP archives containing metadata and weights.
"""

import zipfile
import json
import h5py
import os
import shutil

KERAS_PATH = "./models/aging_score_autoencoder.keras"
WEIGHTS_PATH = "./models/aging_score_autoencoder.weights.h5"
BACKUP_PATH = "./models/aging_score_autoencoder.keras.bak"

print(f"Extracting weights from {KERAS_PATH}...")

if not os.path.exists(KERAS_PATH):
    print(f"ERROR: {KERAS_PATH} not found!")
    exit(1)

try:
    # Backup original
    if not os.path.exists(BACKUP_PATH):
        shutil.copy(KERAS_PATH, BACKUP_PATH)
        print(f"Backed up to {BACKUP_PATH}")
    
    # Extract weights from .keras ZIP
    with zipfile.ZipFile(KERAS_PATH, 'r') as zf:
        # List contents
        files = zf.namelist()
        print(f"Files in {KERAS_PATH}: {files}")
        
        # Find the weights file (usually 'model.weights.h5')
        weights_file = None
        for f in files:
            if 'weights' in f.lower() and f.endswith('.h5'):
                weights_file = f
                break
        
        if weights_file:
            print(f"Found weights file: {weights_file}")
            # Extract to temp location
            with zf.open(weights_file) as src:
                with open(WEIGHTS_PATH, 'wb') as dst:
                    dst.write(src.read())
            print(f"Extracted weights to {WEIGHTS_PATH}")
        else:
            print("ERROR: No .h5 weights file found in archive!")
            print(f"Contents: {files}")
            exit(1)

    print("\nSuccess! Now update app.py to use the weights file.")
    print(f"Change MODEL_PATH to: {WEIGHTS_PATH}")
    
except Exception as e:
    print(f"ERROR: {e}")
    import traceback
    traceback.print_exc()
    exit(1)