#!/usr/bin/env python3
"""Conservative, explicit disk preparation for an inspected Arch live environment.

CONTRACT: no implicit target or safety bypass. JSON inspection failures forbid erasure.
INVARIANT: every destructive command is checked before the next one runs.
"""
import argparse
import fcntl
import json
import os
from pathlib import Path
import re
import shutil
import signal
import subprocess
import sys

MOUNT = Path('/mnt')
SUBVOLS = ('@', '@home', '@log', '@pkg')
OPTS = 'rw,relatime,ssd,discard=async,space_cache=v2'
STEPS = []
MOUNTS = []
CURRENT = 'préconditions'

class Stop(Exception):
    pass

class Interrupted(Stop):
    def __init__(self, number):
        super().__init__(f'signal {number}')
        self.code = 128 + number

def command(args, *, input_text=None):
    try:
        p = subprocess.run(args, input=input_text, text=True, capture_output=True, check=False)
    except OSError as exc:
        raise Stop(f"commande indisponible {args[0]} : {exc}") from exc
    if p.returncode:
        raise Stop(f"{args[0]} a échoué (code {p.returncode}) : {p.stderr.strip()[:200]}")
    return p.stdout

def structured(args, key):
    try:
        data = json.loads(command(args))
    except (ValueError, TypeError) as exc:
        raise Stop(f"réponse structurée invalide de {args[0]}") from exc
    if not isinstance(data, dict) or not isinstance(data.get(key), list):
        raise Stop(f"réponse structurée incomplète de {args[0]}")
    return data[key]

def disk_tree(path):
    rows = structured(['lsblk', '--json', '--tree', '--bytes', '--output',
                       'PATH,TYPE,MAJ:MIN,SIZE,MODEL,SERIAL,WWN,MOUNTPOINTS,PKNAME,PARTN', path], 'blockdevices')
    if len(rows) != 1 or not isinstance(rows[0], dict):
        raise Stop('identité du disque ambiguë')
    return rows[0]

def descendants(node):
    if not isinstance(node, dict) or not re.fullmatch(r'\d+:\d+', str(node.get('maj:min', ''))):
        raise Stop('topologie bloc incomplète')
    yield node
    children = node.get('children', [])
    if not isinstance(children, list):
        raise Stop('descendants invalides')
    for child in children:
        yield from descendants(child)

def mounts():
    rows = structured(['findmnt', '--json', '--output', 'TARGET,SOURCE,MAJ:MIN,FSTYPE,OPTIONS,ID'], 'filesystems')
    def flatten(items):
        for item in items:
            if not isinstance(item, dict) or not isinstance(item.get('target'), str):
                raise Stop('montage illisible')
            yield item
            children = item.get('children', [])
            if not isinstance(children, list):
                raise Stop('arbre de montage invalide')
            yield from flatten(children)
    return list(flatten(rows))

def containing_mount(path, rows):
    path = os.path.realpath(path)
    matches = [m for m in rows if path == m['target'] or path.startswith(m['target'].rstrip('/') + '/')]
    if not matches:
        raise Stop(f'montage porteur inconnu : {path}')
    return max(matches, key=lambda m: len(m['target']))

def target_identity(path):
    # readlink -e rejects a missing path and resolves stable /dev/disk/by-id aliases.
    return command(['readlink', '-e', '--', path]).strip()

def device_source(value):
    # findmnt prints Btrfs subvolumes as /dev/name[/@]; the bracketed suffix
    # describes the mount, not a different block device.
    return os.path.realpath(str(value).split('[', 1)[0])

def mounted_subvolume(options):
    match = re.search(r'(?:^|,)subvol=/?([^,]+)(?:,|$)', str(options))
    return match.group(1) if match else None

def overlay_backing_safe(root_mount, all_mounts, target_ids, boot_source):
    # A visible boot medium alone is insufficient: the overlay may use another
    # disk for a lower or upper layer. Resolve every declared backing path.
    options = {}
    for field in str(root_mount.get('options', '')).split(','):
        if '=' in field:
            key, value = field.split('=', 1)
            options[key] = value
    if not options.get('lowerdir') or not options.get('upperdir') or not options.get('workdir'):
        return False
    paths = options['lowerdir'].split(':') + [options['upperdir'], options['workdir']]
    for path in paths:
        if not path.startswith('/run/archiso/') or not Path(path).exists():
            return False
        carrier = containing_mount(path, all_mounts)
        if carrier['target'] == '/' or carrier.get('maj:min') in target_ids:
            return False
        if carrier.get('fstype') == 'tmpfs':
            continue
        if carrier['target'] == '/run/archiso/bootmnt' and device_source(carrier.get('source', '')) == boot_source:
            continue
        source = device_source(carrier.get('source', ''))
        if not source.startswith('/dev/loop'):
            return False
        backing_file = Path('/sys/block') / Path(source).name / 'loop/backing_file'
        try:
            backing = backing_file.read_text(encoding='utf-8').strip()
        except OSError:
            return False
        if not backing.startswith('/run/archiso/bootmnt/'):
            return False
        if containing_mount(backing, all_mounts)['target'] != '/run/archiso/bootmnt':
            return False
    return True

def inspect(path, *, application):
    resolved = target_identity(path)
    if not resolved.startswith('/dev/') or '\n' in resolved:
        raise Stop('cible hors /dev ou chemin ambigu')
    tree = disk_tree(resolved)
    if tree.get('type') != 'disk' or tree.get('path') != resolved:
        raise Stop('la cible doit être un disque entier pris en charge')
    majmin = str(tree.get('maj:min', ''))
    size = tree.get('size')
    if not re.fullmatch(r'\d+:\d+', majmin) or not isinstance(size, int) or size < 2 * 1024**3:
        raise Stop('identité ou capacité du disque invalide')
    nodes = list(descendants(tree))
    ids = {n['maj:min'] for n in nodes}
    if len(ids) != len(nodes) or any(n.get('type') not in ('disk', 'part') for n in nodes):
        raise Stop('pile bloc non prise en charge')
    for n in nodes:
        if any(x is not None for x in (n.get('mountpoints') or [])):
            raise Stop('disque ou descendant monté selon lsblk')
        holders = Path('/sys/dev/block') / n['maj:min'] / 'holders'
        try:
            if any(holders.iterdir()):
                raise Stop(f'holder actif : {n["maj:min"]}')
        except OSError as exc:
            raise Stop(f'holders illisibles : {exc}') from exc
    all_mounts = mounts()
    for m in all_mounts:
        if m.get('maj:min') in ids:
            raise Stop('disque ou descendant monté selon findmnt')
    # An Arch ISO overlay is acceptable only when its physical boot medium is
    # visible as a separate block mount; all visible mounts are compared above.
    # Other overlays and multi-device Btrfs remain ambiguous and fail closed.
    boot = [m for m in all_mounts if m['target'] == '/run/archiso/bootmnt']
    live_overlay_mapped = False
    if len(boot) == 1 and str(boot[0].get('source', '')).startswith('/dev/'):
        boot_device = disk_tree(str(boot[0]['source']))
        live_overlay_mapped = (boot_device.get('type') in ('disk', 'part') and
                               boot_device.get('maj:min') == boot[0].get('maj:min') and
                               boot_device.get('maj:min') not in ids and
                               boot_device.get('pkname') != Path(resolved).name)
        if live_overlay_mapped:
            root_mount = containing_mount('/', all_mounts)
            if root_mount.get('fstype') == 'overlay':
                live_overlay_mapped = overlay_backing_safe(root_mount, all_mounts, ids,
                                                           device_source(boot[0]['source']))
    for protected in ('/', str(Path(__file__).resolve().parent.parent), '/run/archiso'):
        m = containing_mount(protected, all_mounts)
        if m.get('fstype') == 'overlay' and m['target'] == '/' and live_overlay_mapped:
            continue
        if m.get('fstype') in ('overlay', 'btrfs') and m.get('maj:min') not in ids:
            if application:
                raise Stop(f'stockage protégé non attribuable sans ambiguïté : {protected}')
            print(f'[!] Simulation : provenance non établie pour {protected}; application interdite sans résolution.', file=sys.stderr)
        if m.get('maj:min') in ids:
            raise Stop(f'disque utilisé par {protected}')
    swaps = command(['swapon', '--show=NAME', '--noheadings', '--raw'])
    for line in swaps.splitlines():
        if not line or not line.startswith('/'):
            raise Stop('entrée swap ambiguë')
        real = os.path.realpath(line)
        if real.startswith('/dev/'):
            root = disk_tree(real)
            if root['maj:min'] in ids:
                raise Stop('swap actif sur la cible')
            # Other block stacks are not interpreted as evidence that target is free.
            if root.get('type') not in ('disk', 'part'):
                raise Stop('swap sur pile bloc ambiguë')
        elif containing_mount(real, all_mounts).get('maj:min') in ids:
            raise Stop('swapfile actif sur la cible')
    try:
        if MOUNT.is_symlink() or not MOUNT.is_dir() or any(MOUNT.iterdir()):
            raise Stop('/mnt lié, absent ou non vide')
    except OSError as exc:
        raise Stop(f'/mnt illisible : {exc}') from exc
    if any(m['target'] == '/mnt' or m['target'].startswith('/mnt/') for m in all_mounts):
        raise Stop('/mnt ou descendant déjà monté')
    if application:
        if os.geteuid() != 0 or not Path('/sys/firmware/efi').is_dir() or not Path('/run/archiso').is_dir():
            raise Stop('application réservée au live Arch UEFI avec privilèges root')
    return resolved, majmin, size, tree

def step(label, args):
    global CURRENT
    CURRENT = label
    print(f'[étape] {label}: {" ".join(args)}', flush=True)
    command(args)
    STEPS.append(label)

def mount(label, args, path, source):
    step(label, args)
    found = [m for m in mounts() if m['target'] == path]
    if len(found) != 1 or device_source(found[0].get('source', '')) != source:
        raise Stop(f'effet du montage non établi : {path}')
    mount_id = found[0].get('id')
    if not isinstance(mount_id, int) or mount_id <= 0:
        raise Stop(f'identité du montage indisponible : {path}')
    if '-o' in args and mounted_subvolume(found[0].get('options', '')) != args[2].split('subvol=', 1)[-1]:
        raise Stop(f'option de sous-volume non établie : {path}')
    # Record only the expected mount made by this invocation; never unmount
    # an object that another actor replaced at the same path.
    MOUNTS.append((path, source, mount_id))

def cleanup():
    remaining = []
    for path, source, mount_id in reversed(MOUNTS):
        try:
            found = [m for m in mounts() if m['target'] == path]
            if len(found) != 1 or device_source(found[0].get('source', '')) != source or found[0].get('id') != mount_id:
                remaining.append(path)
                print(f'[!] montage à examiner, non démonté : {path}', file=sys.stderr)
                continue
            command(['umount', '--', path])
        except Stop as exc:
            remaining.append(path)
            print(f'[!] démontage ciblé échoué : {path}: {exc}', file=sys.stderr)
    return remaining

def main():
    parser = argparse.ArgumentParser(description='Préparation destructive GPT/ESP/Btrfs ; application réservée au live Arch UEFI.')
    parser.add_argument('--dry-run', action='store_true', help='plan et inspections lisibles, aucun effet ni confirmation')
    parser.add_argument('disk', help='disque entier explicite, jamais une partition')
    args = parser.parse_args()
    resolved, majmin, size, tree = inspect(args.disk, application=not args.dry_run)
    print(f'Cible : {resolved} ; identité {majmin} ; taille {size} octets ; modèle {tree.get("model") or "inconnu"} ; identifiant stable {tree.get("wwn") or tree.get("serial") or "indisponible"}', flush=True)
    print('Plan : effacer signatures, GPT, ESP FAT32 1 Gio /boot, Btrfs @ @home @log @pkg @/.snapshots sous /mnt.', flush=True)
    if args.dry_run:
        print('Simulation : aucun privilège, confirmation, verrou, montage ou effacement demandé. Les propriétés non lisibles ne sont pas validées.', flush=True)
        return 0
    for tool in ('wipefs','parted','partprobe','udevadm','mkfs.fat','mkfs.btrfs','btrfs','mount','umount','mkdir'):
        if shutil.which(tool) is None:
            raise Stop(f'dépendance absente : {tool}')
    lockpath = f'/run/lock/arch-system-prepare-{majmin.replace(":", "_")}.lock'
    try:
        fd = os.open(lockpath, os.O_CREAT | os.O_RDWR | os.O_NOFOLLOW | os.O_CLOEXEC, 0o600)
        lock = os.fdopen(fd, 'r+', encoding='ascii')
    except OSError as exc:
        raise Stop(f'verrou inaccessible : {exc}') from exc
    with lock:
        try:
            fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError as exc:
            raise Stop('autre invocation coopérante sur cette cible') from exc
        answer = input(f'Taper exactement ERASE {resolved} {majmin} : ')
        if answer != f'ERASE {resolved} {majmin}':
            raise Stop('confirmation refusée')
        again = inspect(args.disk, application=True)
        if again[:3] != (resolved, majmin, size):
            raise Stop('identité changée avant effacement')
        try:
            step('effacement des signatures', ['wipefs','-af','--',resolved])
            step('GPT', ['parted','-s',resolved,'mklabel','gpt'])
            step('ESP', ['parted','-s',resolved,'mkpart','ESP','fat32','1MiB','1025MiB'])
            step('drapeau ESP', ['parted','-s',resolved,'set','1','esp','on'])
            step('partition Btrfs', ['parted','-s',resolved,'mkpart','ARCH','btrfs','1025MiB','100%'])
            step('partprobe', ['partprobe',resolved])
            step('udev', ['udevadm','settle'])
            rows = list(descendants(disk_tree(resolved)))
            parts = [r for r in rows if r.get('type') == 'part']
            if len(parts) != 2 or {str(p.get('partn')) for p in parts} != {'1', '2'} or any(not isinstance(p.get('path'), str) or not p['path'].startswith('/dev/') or p.get('pkname') != Path(resolved).name for p in parts):
                raise Stop('partitions attendues non établies')
            efi = next(p['path'] for p in parts if str(p['partn']) == '1')
            root = next(p['path'] for p in parts if str(p['partn']) == '2')
            if efi == root or any(p.get('maj:min') == majmin for p in parts):
                raise Stop('identité des partitions invalide')
            step('FAT32', ['mkfs.fat','-F','32','-n','EFI',efi])
            step('Btrfs', ['mkfs.btrfs','-f','-L','ARCH',root])
            mount('montage top-level', ['mount',root,str(MOUNT)], str(MOUNT), root)
            for sub in SUBVOLS:
                step('sous-volume '+sub, ['btrfs','subvolume','create',str(MOUNT / sub)])
            # Create the nested snapshot subvolume inside @, before mounting @.
            step('sous-volume @/.snapshots', ['btrfs','subvolume','create',str(MOUNT / '@' / '.snapshots')])
            step('démontage top-level', ['umount','--',str(MOUNT)])
            if any(m['target'] == str(MOUNT) for m in mounts()):
                raise Stop('démontage top-level non établi')
            MOUNTS.clear()
            mount('montage racine', ['mount','-o',OPTS+',subvol=@',root,str(MOUNT)], str(MOUNT), root)
            for rel in ('boot','home','var/log','var/cache/pacman/pkg'):
                step('point de montage '+rel, ['mkdir','-p','--',str(MOUNT / rel)])
            for rel, sub in (('home','@home'),('var/log','@log'),('var/cache/pacman/pkg','@pkg'),('.snapshots','@/.snapshots')):
                mount('montage '+rel, ['mount','-o',OPTS+',subvol='+sub,root,str(MOUNT / rel)], str(MOUNT / rel), root)
            mount('montage ESP', ['mount',efi,str(MOUNT / 'boot')], str(MOUNT / 'boot'), efi)
            final = mounts()
            expected = {str(MOUNT):root, str(MOUNT/'boot'):efi, **{str(MOUNT/rel):root for rel in ('home','var/log','var/cache/pacman/pkg','.snapshots')}}
            for path, source in expected.items():
                found = [m for m in final if m['target'] == path]
                if len(found) != 1 or device_source(found[0].get('source','')) != source:
                    raise Stop('schéma final non vérifié : '+path)
                expected_subvol = {'/mnt':'@','/mnt/home':'@home','/mnt/var/log':'@log',
                                   '/mnt/var/cache/pacman/pkg':'@pkg','/mnt/.snapshots':'@/.snapshots'}.get(path)
                if expected_subvol and mounted_subvolume(found[0].get('options','')) != expected_subvol:
                    raise Stop('sous-volume final non vérifié : '+path)
            print('[+] Disque préparé et montages vérifiés sous /mnt ; Arch non installé, démarrage non testé.', flush=True)
            return 0
        except Stop:
            raise

def entry():
    try:
        return main()
    except (Stop, EOFError, KeyboardInterrupt) as exc:
        print(f'[!] Arrêt : {exc or "entrée interrompue"}; étape à vérifier : {CURRENT}', file=sys.stderr)
        print('[!] Étapes accomplies : '+(', '.join(STEPS) or 'aucune'), file=sys.stderr)
        if MOUNTS:
            cleanup()
        try:
            remaining = [m['target'] for m in mounts() if m['target'] == '/mnt' or m['target'].startswith('/mnt/')]
            print('[!] Montages restants à examiner : '+(', '.join(remaining) or 'aucun'), file=sys.stderr)
        except Stop as read_error:
            print(f'[!] Montages restants non établis : {read_error}', file=sys.stderr)
        print('[!] Ne pas relancer sans inspection du disque et des effets de la dernière étape.', file=sys.stderr)
        return exc.code if isinstance(exc, Interrupted) else (130 if isinstance(exc, KeyboardInterrupt) else 1)

def handle_term(number, frame):
    # Keep the signal's conventional exit code while reporting any mounts
    # that still need human inspection after a partially completed operation.
    raise Interrupted(number)

if __name__ == '__main__':
    signal.signal(signal.SIGTERM, handle_term)
    sys.exit(entry())