#!/usr/bin/env python3
import os
import base64
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives import serialization

def generate_vapid_keys():
    """
    Generate VAPID keys and update .env file, preserving existing variables.
    Only regenerates keys if they are missing or empty.
    """
    # Read existing .env file if it exists
    env_vars = {}
    if os.path.exists('.env'):
        with open('.env', 'r') as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith('#') and '=' in line:
                    key, value = line.split('=', 1)
                    env_vars[key.strip()] = value.strip()

    # Check if we need to generate new keys
    need_new_keys = (
        'VAPID_PRIVATE_KEY' not in env_vars or
        'VAPID_PUBLIC_KEY' not in env_vars or
        not env_vars.get('VAPID_PRIVATE_KEY') or
        not env_vars.get('VAPID_PUBLIC_KEY')
    )

    if need_new_keys:
        # Generate EC private key
        private_key = ec.generate_private_key(ec.SECP256R1())

        # Serialize private key to PEM format
        private_pem = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption()
        ).decode('utf-8')

        # Format for .env file (replace newlines with \n)
        env_private_key = private_pem.strip().replace('\n', '\\n')
        
        # Get public key
        public_key = private_key.public_key()
        public_key_bytes = public_key.public_bytes(
            encoding=serialization.Encoding.X962,
            format=serialization.PublicFormat.UncompressedPoint
        )

        # Store keys
        env_vars['VAPID_PRIVATE_KEY'] = env_private_key  # Single-line format
        env_vars['VAPID_PUBLIC_KEY'] = base64.urlsafe_b64encode(public_key_bytes).decode('utf-8')
        
        print("New VAPID keys generated in .env-compatible format.")
    else:
        print("Existing VAPID keys found - no changes made.")
        # Verify existing key format
        if '-----BEGIN PRIVATE KEY-----' not in env_vars['VAPID_PRIVATE_KEY']:
            print("Warning: Existing private key doesn't appear to be in PEM format!")

    # Ensure we have all required configuration variables with defaults if missing
    defaults = {
        # Flic Button Configuration
        'FLIC_BUTTON1_SERIAL': env_vars.get('FLIC_BUTTON1_SERIAL', 'your_button1_serial'),
        'FLIC_BUTTON2_SERIAL': env_vars.get('FLIC_BUTTON2_SERIAL', 'your_button2_serial'),
        'FLIC_BUTTON3_SERIAL': env_vars.get('FLIC_BUTTON3_SERIAL', 'your_button3_serial'),
        
        # Subscription Storage
        'SUBSCRIPTIONS_FILE': env_vars.get('SUBSCRIPTIONS_FILE', 'data/subscriptions.json'),
        
        # Logging Configuration
        'LOG_LEVEL': env_vars.get('LOG_LEVEL', 'INFO'),
        
        # VAPID Claim (email)
        'VAPID_CLAIM_EMAIL': env_vars.get('VAPID_CLAIM_EMAIL', 'mailto:your-email@example.com')
    }

    # Update env_vars with defaults for any missing keys
    env_vars.update({k: v for k, v in defaults.items() if k not in env_vars})

    # Write back to .env file
    with open('.env', 'w') as f:
        f.write("# VAPID Keys for Web Push\n")
        f.write(f"VAPID_PRIVATE_KEY={env_vars['VAPID_PRIVATE_KEY']}\n")
        f.write(f"VAPID_PUBLIC_KEY={env_vars['VAPID_PUBLIC_KEY']}\n\n")
        
        f.write("# Flic Button Configuration\n")
        f.write(f"FLIC_BUTTON1_SERIAL={env_vars['FLIC_BUTTON1_SERIAL']}\n")
        f.write(f"FLIC_BUTTON2_SERIAL={env_vars['FLIC_BUTTON2_SERIAL']}\n")
        f.write(f"FLIC_BUTTON3_SERIAL={env_vars['FLIC_BUTTON3_SERIAL']}\n\n")
        
        f.write("# Subscription Storage\n")
        f.write(f"SUBSCRIPTIONS_FILE={env_vars['SUBSCRIPTIONS_FILE']}\n\n")
        
        f.write("# Logging Configuration\n")
        f.write(f"LOG_LEVEL={env_vars['LOG_LEVEL']}\n\n")
        
        f.write("# VAPID Claim Email\n")
        f.write(f"VAPID_CLAIM_EMAIL={env_vars['VAPID_CLAIM_EMAIL']}\n")

if __name__ == '__main__':
    generate_vapid_keys()