#!/usr/bin/env python3
"""Create/resume a self-contained HTTP controller deployment. No proxy or public PKI."""
import argparse
import json
import os
from pathlib import Path
import secrets
import shutil
import hashlib
import subprocess
import sys
import urllib.parse


def run(root, *args):
    subprocess.run(['docker', 'compose', '--project-directory', str(root), '-f', str(root / 'compose.json'), *args], check=True)


def file_hash(path):
    digest = hashlib.sha256()
    with path.open('rb') as source:
        for chunk in iter(lambda: source.read(1024 * 1024), b''):
            digest.update(chunk)
    return digest.hexdigest()


def compose(root, *args, **kwargs):
    return subprocess.run(['docker', 'compose', '--project-directory', str(root), '-f', str(root / 'compose.json'), *args], check=True, **kwargs)


def bind_volumes(directory, fresh=False):
    directory = directory.resolve()
    if fresh and directory.exists() and any(directory.iterdir()):
        raise ValueError('restore data directory must be empty')
    volumes = {}
    for name in ('database', 'server-state'):
        path = directory / name
        path.mkdir(parents=True, exist_ok=True)
        if name == 'server-state':
            os.chown(path, 10001, 10001)
            os.chmod(path, 0o750)
        volumes[name] = {'driver': 'local', 'driver_opts': {'type':'none', 'o':'bind', 'device':str(path)}}
    return volumes


def maintain_deployment(args, root, parser):
    archive = args.archive.resolve() if args.archive else None
    if archive is None:
        parser.error('--archive is required for backup, restore and upgrade')
    names = ('compose.json', 'server.toml', 'master.key', 'database.key')
    if args.action in ('backup', 'upgrade'):
        if args.action == 'upgrade' and args.image == 'yoriai-server:local':
            parser.error('upgrade requires an explicit versioned --image or digest')
        archive.mkdir(mode=0o700, parents=True, exist_ok=False)
        # Stop writes before taking the DB and state snapshots. The old version
        # resumes even if creating the backup fails.
        run(root, 'stop', 'server')
        try:
            for name in names:
                shutil.copyfile(root / name, archive / name)
                os.chmod(archive / name, 0o600)
            with (archive / 'database.dump').open('xb') as output:
                compose(root, 'exec', '-T', 'database', 'pg_dump', '-U', 'yoriai', '-d', 'yoriai', '-Fc', stdout=output)
            with (archive / 'server-state.tar.gz').open('xb') as output:
                compose(root, 'run', '--rm', '--no-deps', '-T', '--entrypoint', 'tar', 'server', '-czf', '-', '-C', '/var/lib/yoriai/server', '.', stdout=output)
            hashes = {p.name: file_hash(p) for p in archive.iterdir() if p.is_file()}
            (archive / 'manifest.json').write_text(json.dumps({'format': 1, 'sha256': hashes}, indent=2) + '\n')
        finally:
            run(root, 'up', '-d', '--wait', 'server')
        print(f'Backup written to {archive}. Keep it private: it includes recovery secrets.')
        if args.action == 'backup':
            return
        # Pull before stopping the working controller. Database migrations are
        # forward-only; rollback is by the complete backup, never by an old binary.
        if subprocess.run(['docker', 'image', 'inspect', args.image], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode:
            subprocess.run(['docker', 'pull', args.image], check=True)
        config = json.loads((root / 'compose.json').read_text())
        config['services']['server']['image'] = args.image
        temporary = root / 'compose.next.json'
        temporary.write_text(json.dumps(config, indent=2) + '\n')
        os.chmod(temporary, 0o640)
        temporary.replace(root / 'compose.json')
        run(root, 'up', '-d', '--wait', 'server')
        run(root, 'exec', '-T', 'server', '/usr/bin/yoriai-server', 'doctor')
        print(f'Upgrade started. Restore {archive} into a fresh deployment if rollback is needed.')
        return
    manifest = json.loads((archive / 'manifest.json').read_text())
    if manifest.get('format') != 1:
        parser.error('unsupported backup format')
    required = (*names, 'database.dump', 'server-state.tar.gz')
    for name in required:
        if file_hash(archive / name) != manifest['sha256'].get(name):
            parser.error(f'backup checksum mismatch: {name}')
    if root.exists() and any(root.iterdir()):
        parser.error('restore requires an empty deployment directory')
    config = json.loads((archive / 'compose.json').read_text())
    config['name'] = args.project
    # Never overwrite an existing deployment's volumes, even if its containers
    # are stopped or its original configuration directory was lost.
    for volume in ('database', 'server-state'):
        name = args.project + '_' + volume
        if subprocess.run(['docker', 'volume', 'inspect', name], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode == 0:
            parser.error(f'volume {name} already exists; choose a fresh --project')
    if args.data_directory or any(v.get('driver_opts') for v in config['volumes'].values()):
        config['volumes'] = bind_volumes(args.data_directory or root / 'data', fresh=True)
    root.mkdir(mode=0o750, parents=True, exist_ok=True)
    os.chown(root, 0, 10001)
    for name in names:
        shutil.copyfile(archive / name, root / name)
        os.chown(root / name, 0, 10001)
        os.chmod(root / name, 0o640)
    (root / 'compose.json').write_text(json.dumps(config, indent=2) + '\n')
    run(root, 'up', '-d', '--wait', 'database')
    with (archive / 'database.dump').open('rb') as source:
        compose(root, 'exec', '-T', 'database', 'pg_restore', '-U', 'yoriai', '-d', 'yoriai', '--exit-on-error', '--no-owner', stdin=source)
    with (archive / 'server-state.tar.gz').open('rb') as source:
        # The state volume and archive belong to the service identity. Root
        # with cap_drop=ALL cannot restore ownership/timestamps on its files.
        compose(root, 'run', '--rm', '--no-deps', '-T', '--entrypoint', 'tar', 'server', '-xzf', '-', '-C', '/var/lib/yoriai/server', stdin=source)
    run(root, 'up', '-d', '--wait', 'server')
    run(root, 'exec', '-T', 'server', '/usr/bin/yoriai-server', 'doctor')
    print('Restored configuration, database, master key and controller state. Agent identities are retained.')


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('action', choices=['init', 'start', 'doctor', 'recover-admin', 'backup', 'restore', 'upgrade'])
    parser.add_argument('--directory', type=Path, default=Path('/etc/yoriai/controller'))
    parser.add_argument('--data-directory', type=Path, help='Store database/controller volumes on this host directory')
    parser.add_argument('--archive', type=Path, help='New backup directory, or backup to restore')
    parser.add_argument('--agent-url', help='Worker-reachable URL, e.g. https://192.168.1.10:8443')
    parser.add_argument('--http-port', type=int, default=8080, help='HTTP UI and user API port')
    parser.add_argument('--project', default='yoriai', help='Compose deployment name')
    parser.add_argument('--bundle', type=Path, default=Path(__file__).resolve().parent.parent, help='Extracted verified release bundle; used to build the local image')
    parser.add_argument('--image', default='yoriai-server:local', help='Verified server image/tag or digest')
    args = parser.parse_args()
    if os.geteuid() != 0:
        parser.error('run with sudo; setup creates protected service-owned configuration')
    root = args.directory.resolve()
    if args.action in ('backup', 'restore', 'upgrade'):
        maintain_deployment(args, root, parser)
        return
    if args.action == 'init':
        if (root / 'compose.json').exists():
            print('Existing configuration retained; resuming startup.')
            run(root, 'up', '-d', '--wait', 'database')
            run(root, 'run', '--rm', '--no-deps', 'server', 'bootstrap')
            run(root, 'up', '-d', '--wait', 'server')
            return
        parsed = urllib.parse.urlsplit(args.agent_url or '')
        if parsed.scheme != 'https' or not parsed.hostname or parsed.username or parsed.password or parsed.path not in ('', '/') or parsed.query or parsed.fragment:
            parser.error('--agent-url must name the worker-reachable HTTPS host and port')
        image_exists = subprocess.run(['docker', 'image', 'inspect', args.image], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode == 0
        if not image_exists:
            if args.image == 'yoriai-server:local' and (args.bundle / 'Dockerfile').is_file():
                subprocess.run(['docker', 'build', '-t', args.image, str(args.bundle)], check=True)
            else:
                subprocess.run(['docker', 'pull', args.image], check=True)
        root.mkdir(parents=True, exist_ok=True)
        os.chmod(root, 0o750)
        os.chown(root, 0, 10001)
        if not 1 <= args.http_port <= 65535:
            parser.error('--http-port must be between 1 and 65535')
        password = (root / 'database.key').read_text() if (root / 'database.key').exists() else secrets.token_hex(32)
        master = (root / 'master.key').read_text() if (root / 'master.key').exists() else secrets.token_hex(32)
        files = {
            'master.key': master,
            'database.key': password,
            'server.toml': '\n'.join([
                'listen = "0.0.0.0:8080"', 'agent_listen = "0.0.0.0:8443"',
                'agent_bootstrap = true', f'agent_url = {json.dumps(args.agent_url.rstrip("/"))}',
                f'database_url = "postgres://yoriai:{password}@database/yoriai"',
                'state_root = "/var/lib/yoriai/server"', '[ui]',
                'root = "/usr/share/yoriai/ui"', 'cookies_secure = false',
                '[enrollment]', 'tls_mode = "required"', 'certificate_validity_days = 30',
                f'controller_names = [{json.dumps(parsed.hostname)}]', ''
            ]),
        }
        for name, content in files.items():
            if (root / name).exists():
                if (root / name).read_text() != content:
                    parser.error(f'{name} differs from requested setup; existing configuration retained')
                continue
            fd = os.open(root / name, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o640)
            with os.fdopen(fd, 'w') as file:
                file.write(content)
                file.flush()
                os.fsync(file.fileno())
            os.chown(root / name, 0, 10001)
            os.chmod(root / name, 0o640)
        server = {
            'image': args.image, 'restart': 'unless-stopped', 'user': '10001:10001',
            'environment': {'YORIAI_MASTER_KEY_FILE': '/run/secrets/master.key'},
            'volumes': ['./server.toml:/etc/yoriai/server.toml:ro', './master.key:/run/secrets/master.key:ro', 'server-state:/var/lib/yoriai/server'],
            'ports': [f'{args.http_port}:8080', f'{parsed.port or 443}:8443'],
            'depends_on': {'database': {'condition': 'service_healthy'}},
            'read_only': True, 'tmpfs': ['/tmp'], 'cap_drop': ['ALL'], 'stop_grace_period': '5m',
            'security_opt': ['no-new-privileges:true'],
            'healthcheck': {'test': ['CMD', '/usr/bin/yoriai-server', 'doctor'], 'interval': '5s', 'timeout': '5s', 'retries': 12, 'start_period': '5s'},
        }
        compose = {'name': args.project, 'services': {'server': server, 'database': {
            'image': 'postgres:17-alpine', 'restart': 'unless-stopped',
            'environment': {'POSTGRES_USER': 'yoriai', 'POSTGRES_DB': 'yoriai', 'POSTGRES_PASSWORD_FILE': '/run/secrets/database.key'},
            'volumes': ['database:/var/lib/postgresql/data', './database.key:/run/secrets/database.key:ro'],
            'healthcheck': {'test': ['CMD-SHELL', 'pg_isready -U yoriai -d yoriai'], 'interval': '2s', 'timeout': '3s', 'retries': 30}
        }}, 'volumes': {'database': {}, 'server-state': {}}}
        if args.data_directory:
            compose['volumes'] = bind_volumes(args.data_directory)
        (root / 'compose.json').write_text(json.dumps(compose, indent=2) + '\n')
        os.chmod(root / 'compose.json', 0o640)
        run(root, 'up', '-d', '--wait', 'database')
        run(root, 'run', '--rm', '--no-deps', 'server', 'bootstrap')
        run(root, 'up', '-d', '--wait', 'server')
        print(f'Open http://<server-address>:{args.http_port} and redeem the invitation above. HTTPS proxies are your choice.')
    elif args.action == 'start':
        run(root, 'up', '-d', '--wait', 'database')
        run(root, 'run', '--rm', '--no-deps', 'server', 'bootstrap')
        run(root, 'up', '-d', '--wait', 'server')
    elif args.action == 'recover-admin':
        run(root, 'run', '--rm', '--no-deps', 'server', 'recover-admin')
    else:
        run(root, 'ps')
        run(root, 'exec', '-T', 'server', '/usr/bin/yoriai-server', 'doctor')


if __name__ == '__main__':
    try:
        main()
    except (OSError, ValueError, subprocess.CalledProcessError) as error:
        print(f'Setup failed: {error}', file=sys.stderr)
        sys.exit(1)
