#!/usr/bin/env python3

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

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', 'coreos-installer']
    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_ignition_config(config):
    """Generate Ignition configuration."""
    system = config['system']
    luks = system['luks']
    storage = system['storage']
    
    # Generate the Ignition config
    ignition_config = {
        "ignition": {
            "version": "3.4.0",
            "config": {
                "merge": []
            }
        },
        "storage": {
            "disks": [
                {
                    "device": "/dev/sda",
                    "partitions": [
                        {
                            "label": "boot",
                            "number": 1,
                            "sizeMiB": storage['boot_size_mib'],
                            "typeGuid": "C12A7328-F81F-11D2-BA4B-00A0C93EC93B"
                        },
                        {
                            "label": "root",
                            "number": 2,
                            "sizeMiB": 0,
                            "typeGuid": "21686148-6449-6E6F-744E-656564454649"
                        }
                    ]
                },
                {
                    "device": "/dev/sdb",
                    "partitions": [
                        {
                            "label": "boot2",
                            "number": 1,
                            "sizeMiB": storage['boot_size_mib'],
                            "typeGuid": "C12A7328-F81F-11D2-BA4B-00A0C93EC93B"
                        },
                        {
                            "label": "root2",
                            "number": 2,
                            "sizeMiB": 0,
                            "typeGuid": "21686148-6449-6E6F-744E-656564454649"
                        }
                    ]
                }
            ],
            "luks": [
                {
                    "name": "root",
                    "device": "/dev/disk/by-partlabel/root",
                    "clevis": {
                        "tpm2": True,
                        "tang": [
                            {
                                "url": luks['tang_url'],
                                "thumbprint": luks['tang_thumbprint']
                            }
                        ],
                        "threshold": 1
                    }
                },
                {
                    "name": "root2",
                    "device": "/dev/disk/by-partlabel/root2",
                    "clevis": {
                        "tpm2": True,
                        "tang": [
                            {
                                "url": luks['tang_url'],
                                "thumbprint": luks['tang_thumbprint']
                            }
                        ],
                        "threshold": 1
                    }
                }
            ],
            "filesystems": [
                {
                    "path": "/boot",
                    "device": "/dev/disk/by-partlabel/boot",
                    "format": "vfat",
                    "label": "boot",
                    "wipeFilesystem": True
                },
                {
                    "path": "/boot2",
                    "device": "/dev/disk/by-partlabel/boot2",
                    "format": "vfat",
                    "label": "boot2",
                    "wipeFilesystem": True
                }
            ]
        },
        "systemd": {
            "units": [
                {
                    "name": "kernel-sign.service",
                    "enabled": True,
                    "contents": "[Unit]\nDescription=Sign new kernel images\nAfter=kernel-install.service\n\n[Service]\nType=oneshot\nExecStart=/usr/local/bin/sign-kernel.sh\nRemainAfterExit=yes\n\n[Install]\nWantedBy=multi-user.target"
                }
            ]
        }
    }

    # Add BTRFS filesystems
    for subvol in storage['subvolumes']:
        ignition_config['storage']['filesystems'].append({
            "path": subvol['path'],
            "device": "/dev/disk/by-id/dm-name-root",
            "format": "btrfs",
            "label": "root",
            "wipeFilesystem": subvol['name'] == "@",
            "options": [
                f"subvol={subvol['name']}",
                f"compress={storage['compression']}"
            ]
        })

    return json.dumps(ignition_config, indent=2)

def download_fcos_image(config):
    """Download FCOS image."""
    image = config['image']
    cmd = [
        'coreos-installer', 'download',
        '-s', image['stream'],
        '-p', 'hetzner',
        '-a', image['arch'],
        '-f', 'raw.xz'
    ]
    subprocess.run(cmd, check=True)

def get_image_file():
    """Get the downloaded image file name."""
    files = sorted(Path('.').glob('fedora-coreos-*-hetzner.x86_64.raw.xz'))
    if not files:
        print("Error: No FCOS image found")
        sys.exit(1)
    return files[-1]

def delete_existing_image(image_name):
    """Delete existing image if it exists."""
    cmd = [
        'hcloud', 'image', 'list',
        '--type=snapshot',
        f'--selector=name={image_name}',
        '--output', 'json'
    ]
    result = subprocess.run(cmd, capture_output=True, text=True, check=True)
    images = json.loads(result.stdout)
    
    if images:
        print("Deleting existing image...")
        subprocess.run(['hcloud', 'image', 'delete', str(images[0]['id'])], check=True)

def create_snapshot(config, image_file):
    """Create new snapshot from image."""
    image = config['image']
    cmd = [
        'hcloud-upload-image', 'upload',
        '--architecture', image['hetzner_arch'],
        '--compression', 'xz',
        '--image-path', str(image_file),
        '--name', image['name'],
        '--labels', f'os=fedora-coreos,channel={image["stream"]}',
        '--description', f'Fedora CoreOS ({image["stream"]}, {image["arch"]}) 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()
    
    # Generate Ignition config
    print("Generating Ignition config...")
    ignition_config = generate_ignition_config(config)
    with open('config.ign', 'w') as f:
        f.write(ignition_config)
    
    # Download FCOS image
    print("Downloading FCOS image...")
    download_fcos_image(config)
    
    # Get image file
    image_file = get_image_file()
    
    # Delete existing image if it exists
    print("Checking for existing image...")
    delete_existing_image(config['image']['name'])
    
    # Create new snapshot
    print("Creating snapshot from image...")
    create_snapshot(config, image_file)
    
    print("Image build complete!")

if __name__ == '__main__':
    main()