#!/usr/bin/env python3
# SPDX-License-Identifier: MPL-2.0
"""Create a new private data image; never formats a phone partition."""
import argparse
import json
import os
from pathlib import Path
import shlex
import shutil
import subprocess
import uuid
from persistent_state import FEATURES, digest


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', type=Path, required=True)
    parser.add_argument('--system-image', type=Path, required=True)
    parser.add_argument('--size-mib', type=int, default=2048)
    parser.add_argument('--profile-dir', type=Path,
                        help='Copy an existing droidian-trial browser profile into a fresh P23 state')
    args = parser.parse_args()
    if args.size_mib < 512 or args.size_mib > 8192:
        parser.error('Choose 512–8192 MiB')
    system = json.loads(args.system_image.read_text())
    base = args.system_image.parent/system['image']
    if system.get('device') != 'sargo' or not system.get('filesystemCheckPassed') or digest(base) != system['imageSha256']:
        raise ValueError('System image provenance mismatch')
    os.umask(0o077)
    out = args.output.resolve()
    out.mkdir(parents=True, exist_ok=False)
    seed = out/'seed'
    seed.mkdir()
    (seed/'system.sha256').write_text(system['imageSha256']+'\n')
    for directory in ('upper','work'):
        (seed/directory).mkdir(mode=0o755 if directory == 'upper' else 0o700)
    (seed/'upper').chmod(0o755)
    migrated = []
    opt_layer = None
    if args.profile_dir:
        profile = args.profile_dir.resolve(strict=True)
        if not profile.is_dir() or profile.is_symlink():
            raise ValueError('Profile source must be a real directory')
        opt_layer = seed/'upper/opt'
        vulpes_layer = opt_layer/'vulpes'
        profiles_layer = vulpes_layer/'profiles'
        profiles_layer.mkdir(parents=True)
        opt_layer.chmod(0o755)
        vulpes_layer.chmod(0o755)
        profiles_layer.chmod(0o755)
        target = profiles_layer/'droidian-trial'
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copytree(profile, target, symlinks=True)
        migrated = [vulpes_layer, profiles_layer, target, *target.rglob('*')]
    identifier = str(uuid.uuid4())
    image = out/(identifier+'.ext4')
    with image.open('xb') as stream:
        stream.truncate(args.size_mib*1024**2)
    subprocess.run(['mke2fs','-q','-t','ext4','-F','-b','4096','-I','256','-O',FEATURES,
                    '-U',identifier,'-L','TUNDRA_DATA','-d',str(seed),str(image)],check=True)
    # Root owns the overlay metadata, regardless of the build account's UID.
    for name in ('/system.sha256','/upper','/work'):
        for field in ('uid','gid'):
            subprocess.run(['debugfs','-w','-R',f'set_inode_field {name} {field} 0',str(image)],
                           check=True,stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL)
    if migrated:
        # Gaia runs as the droidian account (UID/GID 32011) on Sargo.
        commands = out/'profile-ownership.debugfs'
        commands.write_text(''.join(
            f'set_inode_field {shlex.quote("/upper/" + str(path.relative_to(seed/"upper")))} {field} 32011\n'
            for path in migrated for field in ('uid','gid')) +
            'set_inode_field /upper/opt uid 0\nset_inode_field /upper/opt gid 0\n')
        subprocess.run(['debugfs','-w','-f',str(commands),str(image)],check=True,
                       stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL)
        commands.unlink()
    check = subprocess.run(['e2fsck','-f','-n',str(image)],text=True,capture_output=True)
    (out/'filesystem-check.txt').write_text(check.stdout+check.stderr)
    if check.returncode:raise RuntimeError('Data image failed filesystem verification')
    manifest = {'schema':1,'device':'sargo','uuid':identifier,'image':image.name,
                'imageSha256':digest(image),'bytes':image.stat().st_size,
                'systemImageSha256':system['imageSha256'],'filesystemCheckPassed':True,
                'migratedProfile':bool(args.profile_dir),
                'scope':'Private persistent overlay; containing filesystem retained; no partition formatting'}
    (out/'state.json').write_text(json.dumps(manifest,indent=2)+'\n')
    print(out/'state.json')


if __name__ == '__main__':main()