#!/usr/bin/env python3

import os
import sys
import yaml
import subprocess
import json
import shutil
from pathlib import Path
import tempfile

def load_config(config_file):
    """Load and parse YAML config file."""
    with open(config_file, 'r') as f:
        return yaml.safe_load(f)

def check_prerequisites():
    """Check if required tools are installed."""
    required_tools = ['hcloud', 'hcloud-upload-image', 'jq', 'virt-customize', 'cloud-localds']
    for tool in required_tools:
        if not shutil.which(tool):
            print(f"Error: {tool} not found. Please install it first.")
            sys.exit(1)

def check_hetzner_token():
    """Check if HCLOUD_TOKEN is set."""
    if not os.environ.get('HCLOUD_TOKEN'):
        print("Error: HCLOUD_TOKEN environment variable not set")
        sys.exit(1)

def generate_cloud_init(config):
    """Generate cloud-init configuration."""
    system = config['system']
    cloud_init = system['cloud_init']
    tpm_config = system['luks']['tpm']
    
    # Generate TPM config
    tpm_json = {
        'pcr_bank': tpm_config['pcr_bank'],
        'pcr_ids': ','.join(map(str, tpm_config['pcr_ids']))
    }
    
    # Base cloud-init config
    cloud_config = {
        'timezone': cloud_init['timezone'],
        'users': cloud_init['users'],
        'package_update': True,
        'package_upgrade': True,
        'packages': cloud_init['packages'],
        'write_files': [
            {
                'path': '/etc/clevis/tang.conf',
                'content': f"URL={system['luks']['tang_url']}\nThumbprint={system['luks']['tang_thumbprint']}\n",
                'permissions': '0644'
            },
            {
                'path': '/etc/clevis/tpm2.conf',
                'content': json.dumps(tpm_json),
                'permissions': '0644'
            },
            {
                'path': '/root/update-tpm-bindings.sh',
                'content': '''#!/bin/bash
# Script to update TPM bindings after firmware updates
# Usage: sudo ./update-tpm-bindings.sh

set -e

# Check if running as root
if [ "$EUID" -ne 0 ]; then 
    echo "Please run as root"
    exit 1
fi

# Check if TPM is available
if ! tpm2_getcap properties-fixed > /dev/null; then
    echo "Error: TPM not available"
    exit 1
fi

# Check if Tang server is accessible
TANG_URL=$(grep URL /etc/clevis/tang.conf | cut -d= -f2)
if ! curl -s -f "$TANG_URL/adv" > /dev/null; then
    echo "Error: Tang server not accessible"
    exit 1
fi

# Get current PCR values
echo "Current PCR values:"
tpm2_pcrread sha256:$(cat /etc/clevis/tpm2.conf | jq -r '.pcr_ids' | tr -d '[]' | tr ',' ' ')

# Ask for confirmation
read -p "Have you updated firmware? Continue with TPM binding update? [y/N] " -n 1 -r
echo
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
    echo "Aborted"
    exit 1
fi

# Update TPM bindings
echo "Updating TPM bindings..."
for dev in /dev/sda2 /dev/sdb2; do
    echo "Processing $dev..."
    
    # Unbind old TPM binding
    clevis luks unbind -d "$dev" -s 1 || true
    
    # Create new TPM binding
    clevis luks bind -d "$dev" tpm2 -c /etc/clevis/tpm2.conf || {
        echo "Error: Failed to bind TPM to $dev"
        exit 1
    }
    
    # Verify Tang binding
    if ! clevis luks list -d "$dev" | grep -q "tang"; then
        echo "Error: Tang binding not found on $dev"
        exit 1
    fi
done

echo "TPM bindings updated successfully!"
echo "Please reboot to verify the changes."
''',
                'permissions': '0700'
            }
        ],
        'runcmd': [
            # Check if passphrase is set
            '[ -n "$LUKS_PASSPHRASE" ] || { echo "Error: LUKS passphrase not found"; exit 1; }',
            # Verify Tang server is accessible
            f'curl -s -f {system["luks"]["tang_url"]}/adv > /dev/null || {{ echo "Tang server not accessible"; exit 1; }}',
            # Verify TPM is available
            'tpm2_getcap properties-fixed > /dev/null || { echo "TPM not available"; exit 1; }',
            # Create filesystem
            'mkfs.btrfs -f -d raid1 -m raid1 /dev/sda2 /dev/sdb2',
            'mount /dev/sda2 /mnt',
            # Create subvolumes
            'btrfs subvolume create /mnt/@boot',
            'btrfs subvolume create /mnt/@home',
            'btrfs subvolume create /mnt/@db',
            'chattr +C /mnt/@db',
            # Setup LUKS with escaped passphrase
            'echo -n "${LUKS_PASSPHRASE}" | tr -d "\n" | cryptsetup luksFormat /dev/sda2 --type luks2 --key-file -',
            'echo -n "${LUKS_PASSPHRASE}" | tr -d "\n" | cryptsetup luksFormat /dev/sdb2 --type luks2 --key-file -',
            # Setup Clevis with error handling
            'clevis luks bind -d /dev/sda2 tpm2 -c /etc/clevis/tpm2.conf || { echo "TPM bind failed"; exit 1; }',
            'clevis luks bind -d /dev/sda2 tang -c /etc/clevis/tang.conf || { echo "Tang bind failed"; exit 1; }',
            'clevis luks bind -d /dev/sdb2 tpm2 -c /etc/clevis/tpm2.conf || { echo "TPM bind failed"; exit 1; }',
            'clevis luks bind -d /dev/sdb2 tang -c /etc/clevis/tang.conf || { echo "Tang bind failed"; exit 1; }',
            # Setup fstab
            'echo "/dev/mapper/root / btrfs compress=zstd 0 0" >> /etc/fstab',
            'echo "/dev/mapper/root /boot btrfs subvol=@boot,compress=zstd 0 0" >> /etc/fstab',
            'echo "/dev/mapper/root /home btrfs subvol=@home,compress=zstd 0 0" >> /etc/fstab',
            'echo "/dev/mapper/root /db btrfs subvol=@db,nodatacow,noatime,compress=zstd 0 0" >> /etc/fstab',
            # Enable services
            'systemctl enable clevis-luks-askpass.service',
            'systemctl enable clevis-luks-askpass.path'
        ]
    }
    
    return yaml.dump(cloud_config)

def download_fedora_image(config):
    """Download Fedora Server cloud image."""
    image = config['image']
    url = f"https://download.fedoraproject.org/pub/fedora/linux/releases/{image['version']}/Cloud/x86_64/images/Fedora-Cloud-Base-{image['version']}-1.6.x86_64.qcow2"
    
    cmd = ['curl', '-L', '-o', 'fedora-server.qcow2', url]
    subprocess.run(cmd, check=True)

def customize_image(config):
    """Customize the Fedora Server image."""
    # Generate cloud-init config
    cloud_init = generate_cloud_init(config)
    
    # Create cloud-init ISO
    with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml') as f:
        f.write(cloud_init)
        f.flush()
        subprocess.run(['cloud-localds', 'cloud-init.iso', f.name], check=True)
    
    # Customize image
    cmd = [
        'virt-customize',
        '-a', 'fedora-server.qcow2',
        '--selinux-relabel',
        '--run-command', 'dnf clean all',
        '--run-command', 'dnf -y update',
        '--run-command', 'dnf -y install cloud-init',
        '--copy-in', 'cloud-init.iso:/var/lib/cloud/seed/nocloud/'
    ]
    subprocess.run(cmd, check=True)

def create_snapshot(config):
    """Create new snapshot from image."""
    image = config['image']
    cmd = [
        'hcloud-upload-image', 'upload',
        '--architecture', image['hetzner_arch'],
        '--compression', 'xz',
        '--image-path', 'fedora-server.qcow2',
        '--name', image['name'],
        '--labels', f'os=fedora-server,version={image["version"]}',
        '--description', f'Fedora Server {image["version"]} for Nullpoint'
    ]
    subprocess.run(cmd, check=True)

def main():
    """Main entry point."""
    # Load config
    if not os.path.exists('build-config.yaml'):
        print("Error: build-config.yaml not found")
        sys.exit(1)
    
    config = load_config('build-config.yaml')
    
    # Check prerequisites
    check_prerequisites()
    check_hetzner_token()
    
    # Download Fedora Server image
    print("Downloading Fedora Server image...")
    download_fedora_image(config)
    
    # Customize image
    print("Customizing image...")
    customize_image(config)
    
    # Create snapshot
    print("Creating snapshot...")
    create_snapshot(config)
    
    print("Image build complete!")

if __name__ == '__main__':
    main()