#!/usr/bin/env python3

import secrets
import string
import os
import base64
from ruamel.yaml import YAML

def generate_password(length=32):
    """Generate a random password."""
    alphabet = string.ascii_letters + string.digits
    return ''.join(secrets.choice(alphabet) for i in range(length))

def generate_api_key():
    """Generate a random 32-byte key and base64 encode it."""
    return base64.b64encode(secrets.token_bytes(32)).decode('utf-8')

SECRETS_FILE = 'secrets.yaml'
# In a real ESPHome project, secrets are often included from a central location
# but for this script, we'll assume it's in the current directory.
# You might need to adjust this path.
secrets_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), SECRETS_FILE)

yaml = YAML()
yaml.preserve_quotes = True
# To prevent line wrapping
yaml.width = 4096

try:
    with open(secrets_path, 'r') as f:
        secrets_data = yaml.load(f)
    if secrets_data is None:
        secrets_data = {}
except FileNotFoundError:
    print(f"Info: '{SECRETS_FILE}' not found. A new file will be created.")
    secrets_data = {}

# Generate new random passwords
new_api_key = generate_api_key()
new_ota_password = generate_password()

# Update the dictionary with the new passwords
if 'api_password' in secrets_data:
    del secrets_data['api_password']
secrets_data['api_key'] = new_api_key
secrets_data['ota_password'] = new_ota_password

# Write the updated dictionary back to the YAML file
with open(secrets_path, 'w') as f:
    yaml.dump(secrets_data, f)

print(f"Successfully updated '{SECRETS_FILE}'.")
print("New values:")
print(f"  api_key: {new_api_key}")
print(f"  ota_password: {new_ota_password}")