#!/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'] + ['mdadm'],  # Add mdadm package
        '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'
            }
        ],
        # Use cloud-init's native disk setup for partitioning only
        'disk_setup': {
            '/dev/sda': {
                'table_type': 'gpt',
                'layout': [
                    [1024, 'boot'],  # 1GB for boot
                    ['auto', 'data']  # Rest for LUKS1
                ],
                'overwrite': True
            },
            '/dev/sdb': {
                'table_type': 'gpt',
                'layout': [
                    [1024, 'boot'],  # 1GB for boot mirror
                    ['auto', 'data']  # Rest for LUKS2
                ],
                'overwrite': True
            }
        },
        '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 boot RAID1
            'mdadm --create --verbose /dev/md0 --level=1 --raid-devices=2 /dev/sda1 /dev/sdb1 --metadata=1.2',
            # Create filesystem on RAID array
            'mkfs.ext4 -L boot /dev/md0',
            # Mount boot RAID
            'mkdir -p /mnt/boot',
            'mount /dev/md0 /mnt/boot',
            # Setup LUKS on data partitions
            '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; }',
            # Open LUKS volumes
            'echo -n "${LUKS_PASSPHRASE}" | tr -d "\n" | cryptsetup luksOpen /dev/sda2 root_a --key-file -',
            'echo -n "${LUKS_PASSPHRASE}" | tr -d "\n" | cryptsetup luksOpen /dev/sdb2 root_b --key-file -',
            # Create BTRFS on decrypted devices
            'mkfs.btrfs -f -d raid1 -m raid1 /dev/mapper/root_a /dev/mapper/root_b',
            'mount /dev/mapper/root_a /mnt',
            # Create subvolumes
            'btrfs subvolume create /mnt/@home',
            'btrfs subvolume create /mnt/@db',
            'chattr +C /mnt/@db',
            # Setup fstab
            'echo "/dev/md0 /boot ext4 defaults 0 0" >> /etc/fstab',
            'echo "/dev/mapper/root_a / btrfs compress=zstd 0 0" >> /etc/fstab',
            'echo "/dev/mapper/root_a /home btrfs subvol=@home,compress=zstd 0 0" >> /etc/fstab',
            'echo "/dev/mapper/root_a /db btrfs subvol=@db,nodatacow,noatime,compress=zstd 0 0" >> /etc/fstab',
            # Save RAID configuration
            'mdadm --detail --scan > /etc/mdadm.conf',
            # 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 netinstall image."""
    image = config['image']
    url = f"https://download.fedoraproject.org/pub/fedora/linux/releases/{image['version']}/Server/x86_64/iso/Fedora-Server-netinst-x86_64-{image['version']}-1.6.iso"
    
    cmd = ['curl', '-L', '-o', 'fedora-server.iso', url]
    subprocess.run(cmd, check=True)

def generate_kickstart(config):
    """Generate kickstart configuration for automated installation."""
    system = config['system']
    
    # TPM update script content
    tpm_update_script = '''#!/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."
'''
    
    kickstart = f"""# Kickstart configuration for Fedora Server
# Generated for Nullpoint

# System language
lang en_US.UTF-8

# Keyboard layouts
keyboard --vckeymap=us --xlayouts='us'

# Network information
network --bootproto=dhcp --device=link --activate

# Root password
rootpw --lock

# System timezone
timezone {system['timezone']} --utc

# Installation type
text

# Wipe all disk
zerombr
clearpart --all --initlabel

# Disk partitioning information
# Boot partitions (5GB each)
part /boot --fstype=btrfs --size=5120 --ondisk=sda
part /boot --fstype=btrfs --size=5120 --ondisk=sdb

# Main data partitions with LUKS
part / --fstype=btrfs --encrypted --cipher=aes-xts-plain64 --luks-version=luks2 --grow --ondisk=sda
part / --fstype=btrfs --encrypted --cipher=aes-xts-plain64 --luks-version=luks2 --grow --ondisk=sdb

# Package source
url --mirrorlist=http://mirrors.fedoraproject.org/mirrorlist?repo=fedora-$releasever&arch=$basearch
repo --name=fedora
repo --name=updates

# Make sure initial-setup runs
firstboot --reconfig

# Package selection
%packages
@^server-product
@system-tools
btrfs-progs
clevis
clevis-luks
clevis-tang
clevis-tpm2
tpm2-tools
tpm2-tss
cryptsetup
systemd
curl
shim-x64
grub2-efi-x64
%end

# Pre-installation script
%pre
# Create TPM and Tang config files
mkdir -p /etc/clevis
cat > /etc/clevis/tang.conf << EOF
URL={system['luks']['tang_url']}
Thumbprint={system['luks']['tang_thumbprint']}
EOF

cat > /etc/clevis/tpm2.conf << EOF
{{
    "pcr_bank": "{system['luks']['tpm']['pcr_bank']}",
    "pcr_ids": "{','.join(map(str, system['luks']['tpm']['pcr_ids']))}"
}}
EOF
%end

# Post-installation script
%post
# https://unix.stackexchange.com/a/351755 for handling TTY in anaconda
printf "\n=== Nullpoint Installation Progress ===\r\n" > /dev/tty1
printf "Press Alt+F3 to view detailed installation logs\r\n" > /dev/tty1
printf "Press Alt+F1 to return to main installation screen\r\n" > /dev/tty1
printf "Current step: Setting up TPM and Tang...\r\n\n" > /dev/tty1
{{
    # Get the LUKS passphrase that was used during installation
    LUKS_PASSPHRASE=$(cat /tmp/luks-passphrase.txt)
    echo "$LUKS_PASSPHRASE" > /root/luks-passphrase.txt
    chmod 600 /root/luks-passphrase.txt

    # Setup Clevis for TPM and Tang
    printf "Configuring Clevis for TPM and Tang...\r\n" > /dev/tty1
    clevis luks bind -d /dev/sda2 tpm2 -c /etc/clevis/tpm2.conf
    clevis luks bind -d /dev/sda2 tang -c /etc/clevis/tang.conf
    clevis luks bind -d /dev/sdb2 tpm2 -c /etc/clevis/tpm2.conf
    clevis luks bind -d /dev/sdb2 tang -c /etc/clevis/tang.conf

    # Get BTRFS UUID (same for all devices in the filesystem)
    BTRFS_UUID=$(blkid -s UUID -o value /dev/mapper/luks-$(blkid -s UUID -o value /dev/sda2))

    # Create subvolumes
    printf "Creating BTRFS subvolumes...\r\n" > /dev/tty1
    # Mount both devices for RAID1
    mount -t btrfs -o raid1 UUID=$BTRFS_UUID /mnt
    btrfs subvolume create /mnt/@root
    btrfs subvolume create /mnt/@home
    btrfs subvolume create /mnt/@db
    chattr +C /mnt/@db

    # Setup fstab
    printf "Configuring system mount points...\r\n" > /dev/tty1
    cat > /etc/fstab << EOF
UUID=$BTRFS_UUID / btrfs subvol=@root,compress=zstd,raid1 0 0
UUID=$BTRFS_UUID /home btrfs subvol=@home,compress=zstd,raid1 0 0
UUID=$BTRFS_UUID /db btrfs subvol=@db,nodatacow,noatime,compress=zstd,raid1 0 0
EOF

    # Enable services
    printf "Enabling system services...\r\n" > /dev/tty1
    systemctl enable clevis-luks-askpass.service
    systemctl enable clevis-luks-askpass.path

    # Create TPM update script
    printf "Creating TPM update script...\r\n" > /dev/tty1
    cat > /root/update-tpm-bindings.sh << 'EOF'
{tpm_update_script}
EOF

    chmod +x /root/update-tpm-bindings.sh
    printf "\nInstallation complete! The system will reboot shortly.\r\n" > /dev/tty1
    printf "IMPORTANT: LUKS passphrase has been saved to /root/luks-passphrase.txt\r\n" > /dev/tty1
}} 2>&1 | tee -a /root/postinstall.log > /dev/tty3
%end
"""
    return kickstart

def customize_image(config):
    """Customize the Fedora Server image."""
    # Generate kickstart config
    kickstart = generate_kickstart(config)
    
    # Create kickstart file
    with tempfile.NamedTemporaryFile(mode='w', suffix='.ks') as f:
        f.write(kickstart)
        f.flush()
        
        # Create custom ISO with kickstart
        cmd = [
            'mkisofs',
            '-o', 'fedora-server-custom.iso',
            '-b', 'isolinux/isolinux.bin',
            '-c', 'isolinux/boot.cat',
            '-boot-info-table',
            '-no-emul-boot',
            '-boot-load-size', '4',
            '-R',
            '-J',
            '-v',
            '-T',
            '-V', 'Fedora-S-custom',
            '-A', 'Fedora-S-custom',
            'fedora-server.iso',
            f.name
        ]
        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-custom.iso',
        '--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()