"""Prepare only exported Gentoo payloads and configure the dedicated guest network."""
import argparse
import bz2
import io
import runpy
import hashlib
import ipaddress
import json
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys

LABEL = 'HYDRA_OS'
BRIDGE = 'hydrabr0'
SUBNET = ipaddress.ip_network('10.203.0.0/24')
CONNECTION_UUID = 'c84844bd-2955-4a61-9463-5e540140f9a2'


class ImageError(RuntimeError):
    pass


def root_path(root, relative, follow_leaf=True):
    root = Path(root).resolve()
    if root == Path('/') or not root.is_dir():
        raise ImageError('Expected an isolated exported rootfs directory')
    relative = Path(relative)
    if relative.is_absolute() or '..' in relative.parts:
        raise ImageError('Payload path must be relative and remain inside the export')
    path = root / relative
    checked = path if follow_leaf else path.parent
    if not checked.resolve().is_relative_to(root):
        raise ImageError(f'Payload symlink escapes the export: {relative}')
    return path


def remove(root, relative):
    path = root_path(root, relative, follow_leaf=False)
    if path.is_symlink() or path.is_file():
        path.unlink()
    elif path.is_dir():
        shutil.rmtree(path)


def write(root, relative, content, mode=0o644):
    path = root_path(root, relative)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(content)
    path.chmod(mode)
    return path


def sanitize_saved_environments(root):
    """Use Portage's parser to scrub exported declarations without evaluating them."""
    database = root_path(root, 'var/db/pkg')
    paths = sorted(database.glob('*/*/environment.bz2'))
    result = {'files_scanned': len(paths), 'filtered_files': 0, 'variable_names': []}
    if not paths:
        return result
    scripts = sorted(root_path(root, 'usr/lib/portage').glob('*/filter-bash-environment.py'))
    if not scripts:
        raise ImageError('Installed Portage environment filter is required for export')
    script = root_path(root, scripts[0].relative_to(root))
    api = runpy.run_path(str(script))
    pattern = re.compile(rb'(?i:(?:HTTP|HTTPS|FTP|ALL|NO)_PROXY|.*(?:PASSWORD|PASSWD|SECRET|TOKEN|API_KEY|ACCESS_KEY|CREDENTIAL|AUTHORIZATION).*)')
    declarations = re.compile(rb'^(?:declare|export|readonly)\s+(?:-[^ ]+\s+)*([A-Za-z_][A-Za-z0-9_]*)=', re.M)
    names = set()
    for path in paths:
        path = root_path(root, path.relative_to(root))
        try:
            original = bz2.decompress(path.read_bytes())
        except (OSError, ValueError) as error:
            raise ImageError('Cannot read an exported package environment') from error
        filtered = io.BytesIO()
        api['filter_bash_environment'](pattern, io.BytesIO(original), filtered)
        clean = filtered.getvalue()
        if clean != original:
            names.update(name.decode('ascii') for name in declarations.findall(original)
                         if pattern.fullmatch(name))
            path.write_bytes(bz2.compress(clean))
            result['filtered_files'] += 1
    result['variable_names'] = sorted(names)
    return result


def sanitize_rootfs(root):
    """Erase builder identity/secrets in the exported copy, retaining Gentoo updates."""
    for relative in ('etc/portage/gnupg', 'etc/portage/gnupg.hydra-untrusted',
                     'var/lib/hydra-build', 'root', 'home', 'var/cache', 'var/tmp',
                     'tmp', 'run', 'var/log', 'var/lib/systemd/random-seed',
                     'var/lib/NetworkManager', 'etc/NetworkManager/system-connections',
                     'hydra-source', 'hydra-cache', 'hydra-output'):
        remove(root, relative)
    for name in root_path(root, 'etc/ssh').glob('ssh_host_*'):
        remove(root, str(name.relative_to(root)))
    for relative, mode in [('root', 0o700), ('home', 0o755), ('tmp', 0o1777),
                           ('var/tmp', 0o1777), ('run', 0o755), ('var/log', 0o755),
                           ('var/cache', 0o755), ('dev', 0o755), ('proc', 0o755), ('sys', 0o755)]:
        path = root_path(root, relative)
        path.mkdir(parents=True, exist_ok=True)
        path.chmod(mode)
    remove(root, 'etc/machine-id')
    write(root, 'etc/machine-id', '')
    remove(root, 'var/lib/dbus/machine-id')
    path = root_path(root, 'var/lib/dbus/machine-id', follow_leaf=False)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.symlink_to('/etc/machine-id')
    write(root, 'etc/hostname', 'hydra-os\n')
    write(root, 'etc/hosts', '127.0.0.1 localhost\n::1 localhost\n127.0.1.1 hydra-os\n')
    remove(root, 'etc/resolv.conf')
    root_path(root, 'etc/resolv.conf', follow_leaf=False).symlink_to('/run/NetworkManager/resolv.conf')
    audit = sanitize_saved_environments(root)
    write(root, 'usr/share/hydra/environment-sanitization.json', json.dumps(audit, indent=2) + '\n')


def copy_overlay(root, overlay):
    for source in sorted(Path(overlay).rglob('*')):
        destination = root_path(root, source.relative_to(overlay))
        if source.is_symlink():
            raise ImageError('Source overlay must contain ordinary directories/files')
        if source.is_dir():
            destination.mkdir(parents=True, exist_ok=True)
        elif source.is_file():
            destination.parent.mkdir(parents=True, exist_ok=True)
            shutil.copyfile(source, destination)
            destination.chmod(source.stat().st_mode & 0o777)


def _kernel_version(root):
    modules = root_path(root, 'lib/modules')
    versions = sorted(path.name for path in modules.iterdir() if path.is_dir()) if modules.is_dir() else []
    if len(versions) != 1 or not re.fullmatch(r'[A-Za-z0-9_.+-]+', versions[0]):
        raise ImageError('Exactly one installed payload kernel is required')
    return versions[0]


def kernel_image(root):
    version = _kernel_version(root)
    candidates = [root_path(root, 'boot/' + prefix + version)
                  for prefix in ('kernel-', 'vmlinuz-')]
    images = [path for path in candidates if path.is_file()]
    if len(images) != 1:
        raise ImageError('Installed modules require exactly one matching boot kernel image')
    return images[0]


def kernel_version(root):
    kernel_image(root)
    return _kernel_version(root)


def live_kernel_args():
    return [f'root=live:CDLABEL={LABEL}', 'rd.live.image', 'rd.live.overlay.overlayfs=1',
            'ro', 'console=tty0', 'console=ttyS0,115200n8']


def source_files(source):
    source = Path(source)
    result = {}
    for directory in ('image', 'hydra', 'bin', 'session', 'patches'):
        for path in sorted((source / directory).rglob('*')):
            if path.is_file() and '__pycache__' not in path.parts and path.suffix != '.pyc':
                result[str(path.relative_to(source))] = digest(path)
    return result


def source_identity(source):
    return hashlib.sha256(json.dumps(source_files(source), sort_keys=True).encode()).hexdigest()


def payload_manifest(root, source):
    packages = []
    database = root_path(root, 'var/db/pkg')
    for use in sorted(database.glob('*/*/USE')):
        package = use.parent
        built = package / 'BUILD_TIME'
        packages.append({'package': str(package.relative_to(database)), 'use': use.read_text().split(),
                         'build_time': built.read_text().strip() if built.is_file() else None})
    commit = os.environ.get('HYDRA_SOURCE_COMMIT')
    if not commit:
        git = subprocess.run(['git', '-C', str(source), 'rev-parse', 'HEAD'], text=True, capture_output=True)
        commit = git.stdout.strip() if git.returncode == 0 else None
    return {'format': 1, 'label': LABEL, 'kernel': kernel_version(root),
            'live_kernel_args': live_kernel_args(), 'source_commit': commit,
            'source_files': source_files(source), 'sources': json.loads((Path(source) / 'image/sources.json').read_text()),
            'packages': packages}


def artifacts_valid(directory, state):
    artifacts = state.get('artifacts')
    if not isinstance(artifacts, dict) or not artifacts:
        return False
    for name, expected in artifacts.items():
        if not isinstance(name, str) or Path(name).name != name or not isinstance(expected, str):
            return False
        path = Path(directory) / name
        try:
            if path.is_symlink() or not path.is_file() or digest(path) != expected:
                return False
        except OSError:
            return False
    return True


def validate_package_inventory(data):
    if not isinstance(data, list) or not data:
        raise ImageError('Package inventory must be a nonempty array')
    seen = set()
    for record in data:
        if not isinstance(record, dict) or set(record) != {'package', 'use', 'build_time'}:
            raise ImageError('Malformed package inventory entry')
        package = record['package']
        if (not isinstance(package, str) or
            not re.fullmatch(r'[A-Za-z0-9+_.-]+/[A-Za-z0-9+_.-]+-[0-9][A-Za-z0-9+_.-]*', package) or
            package in seen):
            raise ImageError('Invalid or duplicate package inventory identity')
        seen.add(package)
        flags = record['use']
        if (not isinstance(flags, list) or any(not isinstance(flag, str) or
            not re.fullmatch(r'[A-Za-z0-9+_.@-]+', flag) for flag in flags) or len(set(flags)) != len(flags)):
            raise ImageError('Invalid package inventory USE flags')
        built = record['build_time']
        if built is not None and (not isinstance(built, str) or not re.fullmatch(r'[0-9]+', built)):
            raise ImageError('Invalid package inventory build time')
    return data


def read_package_inventory(path, expected_digest=None):
    path = Path(path)
    if path.is_symlink() or not path.is_file():
        raise ImageError('Package inventory must be an ordinary file')
    if expected_digest is not None and digest(path) != expected_digest:
        raise ImageError('Package inventory checksum differs from checkpoint')
    try:
        return validate_package_inventory(json.loads(path.read_text()))
    except (ValueError, UnicodeError) as error:
        raise ImageError('Package inventory is not valid JSON') from error


def enable_user_audio(root):
    root_path(root, 'etc/systemd/user')
    result = subprocess.run(['systemctl', '--root=' + str(root), '--global', 'enable',
                             'pipewire.socket', 'pipewire-pulse.socket', 'wireplumber.service'],
                            text=True, capture_output=True)
    if result.returncode:
        raise ImageError('Cannot enable user audio sockets: ' + result.stderr.strip())


def configure_rootfs(root, source):
    root, source = Path(root), Path(source)
    # Refuse incomplete integrations before changing the exported copy.
    for name in ('bin/hydra', 'bin/hydra-install', 'hydra/session.py', 'hydra/installer.py'):
        if not (source / name).is_file():
            raise ImageError(f'Required committed runtime/installer asset is missing: {name}')
    version = kernel_version(root)
    sanitize_rootfs(root)
    copy_overlay(root, source / 'image/rootfs')
    write(root, 'etc/hydra/live-mode/dracut.conf', (source / 'image/dracut.conf').read_text())
    root_path(root, 'etc/hydra/live-mode/dracut.conf.d').mkdir()
    package = root_path(root, 'usr/lib/hydra/hydra')
    package.mkdir(parents=True, exist_ok=True)
    for module in sorted((source / 'hydra').glob('*.py')):
        shutil.copyfile(module, root_path(root, 'usr/lib/hydra/hydra/' + module.name))
    # Session helpers derive this checkout-compatible launcher from __file__.
    for name in ('hydra', 'hydra-install'):
        write(root, 'usr/lib/hydra/bin/' + name, (source / 'bin' / name).read_text(), 0o755)
    for name, module in (('hydra', 'cli'), ('hydra-install', 'installer')):
        write(root, 'usr/bin/' + name,
              '#!/usr/bin/python3\nimport sys\nsys.path.insert(0, "/usr/lib/hydra")\n'
              f'from hydra.{module} import main\nraise SystemExit(main())\n', 0o755)
    write(root, 'usr/bin/hydra-session', '#!/bin/sh\nexec /usr/bin/hydra session "$@"\n', 0o755)
    write(root, 'usr/bin/hydra-network', '#!/usr/bin/python3\nimport sys\nsys.path.insert(0, "/usr/lib/hydra")\n'
          'from hydra.image import network_main\nraise SystemExit(network_main())\n', 0o755)
    write(root, 'usr/share/hydra/build-manifest.json', json.dumps(payload_manifest(root, source), indent=2) + '\n')
    make = root_path(root, 'etc/portage/make.conf')
    if make.is_file():
        settings = make.read_text()
        settings = re.sub(r'(?<=[\s"\'])-(?:usersandbox|sandbox|network-sandbox|ipc-sandbox|pid-sandbox)(?=[\s"\']|$)', '', settings)
        make.write_text(settings)
    for path in root_path(root, 'etc/sudoers.d').iterdir():
        path.chmod(0o440)
    return version


def check_network_conflicts(routes, addresses):
    try:
        if not isinstance(routes, list) or not isinstance(addresses, list):
            raise ValueError('Expected JSON arrays')
        for route in routes:
            dst = route.get('dst')
            if dst in ('default', None):
                continue
            network = ipaddress.ip_network(dst, strict=False)
            if network.version != 4:
                raise ValueError('Unexpected non-IPv4 route')
            if network.overlaps(SUBNET) and not (route.get('dev') == BRIDGE and network.subnet_of(SUBNET)):
                raise ImageError(f'Guest subnet {SUBNET} conflicts with route {dst} on {route.get("dev")}')
        for interface in addresses:
            for addr in interface['addr_info']:
                if addr.get('family') != 'inet':
                    continue
                network = ipaddress.ip_network(f'{addr["local"]}/{addr["prefixlen"]}', strict=False)
                owned = interface['ifname'] == BRIDGE and addr['local'] == '10.203.0.1' and addr['prefixlen'] == 24
                if network.overlaps(SUBNET) and not owned:
                    raise ImageError(f'Guest subnet {SUBNET} conflicts with address on {interface["ifname"]}')
    except (ValueError, KeyError, TypeError, AttributeError) as error:
        raise ImageError('Cannot safely interpret IPv4 routes/addresses') from error


def configure_network(run=subprocess.run):
    def command(argv, **kwargs):
        result = run(argv, text=True, capture_output=True, **kwargs)
        if result.returncode:
            raise ImageError(f'Network command failed: {argv[0]}: {result.stderr.strip()}')
        return result.stdout
    routes = json.loads(command(['ip', '-json', '-4', 'route', 'show', 'table', 'all']))
    addresses = json.loads(command(['ip', '-json', '-4', 'address', 'show']))
    check_network_conflicts(routes, addresses)
    probe = run(['nmcli', '--get-values', 'connection.uuid,connection.type,connection.interface-name,ipv4.method,ipv4.addresses',
                 'connection', 'show', 'uuid', CONNECTION_UUID], text=True, capture_output=True)
    exists = probe.returncode == 0
    if exists and probe.stdout.strip().splitlines() != [CONNECTION_UUID, 'bridge', BRIDGE, 'manual', '10.203.0.1/24']:
        raise ImageError('Hydra bridge connection exists with unexpected configuration')
    if not exists and probe.returncode != 10:
        raise ImageError('Cannot query NetworkManager connections')
    links = json.loads(command(['ip', '-json', '-details', 'link', 'show']))
    for link in links:
        if link.get('ifname') == BRIDGE and (not exists or link.get('linkinfo', {}).get('info_kind') != 'bridge'):
            raise ImageError('Refusing to replace an existing interface named hydrabr0')
    if not exists:
        command(['nmcli', 'connection', 'add', 'type', 'bridge', 'con-name', 'hydra-guest',
                 'ifname', BRIDGE, 'connection.uuid', CONNECTION_UUID,
                 'connection.autoconnect', 'no', 'bridge.stp', 'no',
                 'ipv4.method', 'manual', 'ipv4.addresses', '10.203.0.1/24',
                 'ipv4.never-default', 'yes', 'ipv6.method', 'disabled'])
    # Activation is always gated by the fresh route/address check, including restarts.
    command(['nmcli', 'connection', 'up', 'uuid', CONNECTION_UUID])
    command(['sysctl', '-w', 'net.ipv4.ip_forward=1'])
    table = run(['nft', 'list', 'table', 'ip', 'hydra_nat'], text=True, capture_output=True)
    rules = ('delete table ip hydra_nat\n' if table.returncode == 0 else '') + '''table ip hydra_nat {
  chain postrouting {
    type nat hook postrouting priority srcnat; policy accept;
    ip saddr 10.203.0.0/24 oifname != "hydrabr0" masquerade
  }
}
'''
    command(['nft', '-f', '-'], input=rules)


def network_main():
    try:
        configure_network()
    except (ImageError, OSError, ValueError) as error:
        print(f'hydra-network: {error}', file=sys.stderr)
        return 1
    return 0


def digest(path):
    value = hashlib.sha256()
    with Path(path).open('rb') as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b''):
            value.update(block)
    return value.hexdigest()


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('step', choices=('configure', 'kernel', 'kernel-image', 'network', 'audio'))
    parser.add_argument('root', nargs='?', type=Path)
    parser.add_argument('--source', type=Path)
    options = parser.parse_args(argv)
    try:
        if options.step == 'network':
            return network_main()
        if options.root is None:
            raise ImageError('An isolated rootfs directory is required')
        if options.step == 'kernel':
            print(kernel_version(options.root))
        elif options.step == 'kernel-image':
            print(kernel_image(options.root))
        elif options.step == 'audio':
            enable_user_audio(options.root)
        else:
            if options.source is None:
                raise ImageError('Source checkout is required')
            configure_rootfs(options.root, options.source)
    except (ImageError, OSError, ValueError) as error:
        print(f'hydra-image: {error}', file=sys.stderr)
        return 1
    return 0


if __name__ == '__main__':
    raise SystemExit(main())