#!/usr/bin/env python3
"""DPL append-only batch coordinator. Adapters must probe before and after writes."""
import argparse
import concurrent.futures
import contextlib
import fcntl
import hashlib
import ipaddress
import importlib.util
import json
import os
from pathlib import Path
import re
import signal
import sqlite3
import subprocess
import sys
import threading
import time

PHASES = ('stage', 'dns', 'provision', 'relay', 'verify', 'publish', 'audit')
SHARED = {'stage': 'panel', 'dns': 'dns', 'relay': 'relay', 'publish': 'panel'}
STOP = threading.Event()
EVENTS = False


def digest(value):
    return hashlib.sha256(json.dumps(value, sort_keys=True, ensure_ascii=False,
                                     separators=(',', ':')).encode()).hexdigest()


def read_json(path):
    return json.loads(Path(path).read_text())


def atomic_json(path, value):
    path = Path(path)
    tmp = path.with_name(path.name + '.tmp-' + str(os.getpid()) + '-' + str(threading.get_ident()))
    with tmp.open('x') as f:
        os.chmod(tmp, 0o600)
        json.dump(value, f, ensure_ascii=False, sort_keys=True)
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, path)


def validate(plan):
    if plan.get('version') != 1 or plan.get('mode') != 'append':
        raise ValueError('Only version=1, mode=append is supported')
    if not re.fullmatch(r'[a-z0-9-]{1,64}', plan.get('batch_id', '')):
        raise ValueError('Invalid batch_id')
    context = plan.get('context', {})
    required = ('panel_vps', 'gateway_vps', 'jump_vps', 'app_root', 'panel_url',
                'subscription_url', 'provision_commit', 'bundle', 'client_binary')
    if any(not context.get(k) for k in required):
        raise ValueError('Missing explicit context fields')
    if not re.fullmatch(r'[0-9a-f]{40}', context['provision_commit']):
        raise ValueError('A full reviewed, pushed Git commit is required')
    for k in ('panel_url', 'subscription_url'):
        if not re.fullmatch(r'https://[a-z0-9.-]+', context[k]):
            raise ValueError('HTTPS origin without token/path required: ' + k)
    adapter = plan.get('adapter')
    if not isinstance(adapter, list) or not adapter or not all(isinstance(x, str) for x in adapter):
        raise ValueError('adapter must be an argv array, never a shell string')
    if not plan.get('nodes') or not isinstance(plan['nodes'], list):
        raise ValueError('A nonempty node list is required')
    seen = {k: set() for k in ('key', 'name', 'vps_id', 'node_address', 'tls_domain', 'relay', 'socks_port', 'number')}
    for n in plan['nodes']:
        if any(k not in n for k in ('key', 'country', 'sequence', 'flag', 'country_name', 'name',
                                    'vps_id', 'node_address', 'listen_port', 'relay_host',
                                    'relay_port', 'tls_domain', 'formal_group_id',
                                    'expected_exit', 'socks_port')):
            raise ValueError('Incomplete node specification')
        if not re.fullmatch(r'[a-z0-9-]{1,48}', n['key']):
            raise ValueError('Invalid deployment key')
        if not re.fullmatch(r'[A-Z]{2}', n['country']):
            raise ValueError('Country must be a two-letter code')
        if type(n['sequence']) is not int or not 1 <= n['sequence'] <= 99:
            raise ValueError('Display sequence must be 01..99')
        expected = f"{n['flag']} {n['country_name']} {n['country']} - {n['sequence']:02d}"
        flag = ''.join(chr(0x1F1E6 + ord(x) - ord('A')) for x in n['country'])
        if n['name'] != expected or n['flag'] != flag:
            raise ValueError('Display name/flag does not match country and sequence')
        for k in ('listen_port', 'relay_port', 'socks_port'):
            if type(n[k]) is not int or not 1024 <= n[k] <= 65535:
                raise ValueError('Invalid port: ' + k)
        if type(n['formal_group_id']) is not int or n['formal_group_id'] < 1:
            raise ValueError('An existing formal group is required')
        ipaddress.ip_address(n['node_address'])
        exit_address=ipaddress.ip_address(n['expected_exit'])
        if getattr(exit_address,'scope_id',None):raise ValueError('Scoped/link-local exit addresses are not supported')
        for k in ('relay_host', 'tls_domain'):
            if not re.fullmatch(r'[a-z0-9](?:[a-z0-9.-]*[a-z0-9])?\.[a-z]{2,}', n[k]):
                raise ValueError('Invalid DNS name: ' + k)
        for k in ('vps_id',):
            if not re.fullmatch(r'VPS-[0-9]{3}', n[k]):
                raise ValueError('Use the exact CoDocs VPS ID')
        identities = {k: n[k] for k in ('key', 'name', 'vps_id', 'node_address', 'tls_domain', 'socks_port')}
        identities.update(relay=(n['relay_host'], n['relay_port']), number=(n['country'], n['sequence']))
        for k, value in identities.items():
            if value in seen[k]:
                raise ValueError('Duplicate batch identity: ' + k)
            seen[k].add(value)
    # The plan carries references, never secret values. Adapters load credentials themselves.
    def check(value):
        if isinstance(value, dict):
            for k, v in value.items():
                if re.search(r'password|secret|token|private_key', k, re.I):
                    raise ValueError('Secret fields are not allowed in a batch plan')
                check(v)
        elif isinstance(value, list):
            for v in value:
                check(v)
    check(plan)
    return plan


@contextlib.contextmanager
def lock(root, name):
    path = root / ('lock-' + hashlib.sha256(name.encode()).hexdigest())
    with path.open('a+') as f:
        os.chmod(path, 0o600)
        while True:
            try:
                fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
                break
            except BlockingIOError:
                if STOP.wait(.1):
                    raise RuntimeError('Stopped while waiting for resource lock')
        try:
            yield
        finally:
            fcntl.flock(f, fcntl.LOCK_UN)


class Ledger:
    def __init__(self, root):
        self.root = Path(root).resolve()
        if self.root == Path('/'):
            raise ValueError('Invalid state root')
        self.root.mkdir(parents=True, mode=0o700, exist_ok=True)
        os.chmod(self.root, 0o700)
        self.path = self.root / 'registry.sqlite3'
        with self.db() as c:
            c.executescript('''
                CREATE TABLE IF NOT EXISTS batches(id TEXT PRIMARY KEY, digest TEXT NOT NULL);
                CREATE TABLE IF NOT EXISTS claims(namespace TEXT, kind TEXT, value TEXT,
                    owner TEXT NOT NULL, PRIMARY KEY(namespace,kind,value));
                CREATE TABLE IF NOT EXISTS evidence(batch TEXT,node TEXT,phase TEXT,
                    dependency TEXT,expires REAL,request TEXT,PRIMARY KEY(batch,node,phase));
                CREATE TABLE IF NOT EXISTS phases(batch TEXT, node TEXT, phase TEXT,
                    status TEXT NOT NULL, attempts INTEGER NOT NULL DEFAULT 0,
                    updated REAL NOT NULL, PRIMARY KEY(batch,node,phase));
            ''')
        os.chmod(self.path, 0o600)

    @contextlib.contextmanager
    def db(self):
        c = sqlite3.connect(self.path, timeout=30)
        try:
            c.execute('PRAGMA synchronous=FULL')
            with c: yield c
        finally: c.close()

    def reserve(self, plan):
        namespace = plan['context']['panel_vps'] + ':' + plan['context']['app_root']
        with self.db() as c:
            c.execute('BEGIN IMMEDIATE')
            prior = c.execute('SELECT digest FROM batches WHERE id=?', (plan['batch_id'],)).fetchone()
            if prior and prior[0] != digest(plan):
                raise ValueError('Batch plan changed; resume requires the original exact plan')
            for n in plan['nodes']:
                owner = digest({'context': plan['context'], 'node': n})
                claims = {k: n[k] for k in ('key', 'name', 'vps_id', 'node_address', 'tls_domain')}
                claims['relay'] = [n['relay_host'], n['relay_port']]
                claims['number'] = [n['country'], n['sequence']]
                for kind, value in claims.items():
                    scope = namespace
                    if kind == 'relay':
                        scope = 'relay:' + n['relay_host']
                    elif kind in ('vps_id', 'node_address', 'tls_domain'):
                        scope = 'infrastructure'
                    encoded = json.dumps(value, ensure_ascii=False, sort_keys=True)
                    existing = c.execute('SELECT owner FROM claims WHERE namespace=? AND kind=? AND value=?',
                                         (scope, kind, encoded)).fetchone()
                    if existing and existing[0] != owner:
                        raise ValueError('Another batch reserved identity: ' + kind)
                    c.execute('INSERT OR IGNORE INTO claims VALUES(?,?,?,?)', (scope, kind, encoded, owner))
                # SOCKS listeners are shared by all batches using this gateway.
                scope, value = 'gateway:' + plan['context']['gateway_vps'], str(n['socks_port'])
                old = c.execute('SELECT owner FROM claims WHERE namespace=? AND kind=? AND value=?',
                                (scope, 'socks_port', value)).fetchone()
                if old and old[0] != owner:
                    raise ValueError('Another batch reserved the gateway SOCKS port')
                c.execute('INSERT OR IGNORE INTO claims VALUES(?,?,?,?)', (scope, 'socks_port', value, owner))
            formal_port=plan['context'].get('formal_socks_port')
            if formal_port is not None:
                if type(formal_port) is not int or not 1024<=formal_port<=65535 or any(n['socks_port']==formal_port for n in plan['nodes']):
                    raise ValueError('A distinct formal SOCKS port is required')
                scope,value='gateway:'+plan['context']['gateway_vps'],str(formal_port)
                owner=digest({'context':plan['context'],'formal_port':formal_port})
                old=c.execute('SELECT owner FROM claims WHERE namespace=? AND kind=? AND value=?',(scope,'socks_port',value)).fetchone()
                if old and old[0]!=owner:raise ValueError('Formal SOCKS port already reserved')
                c.execute('INSERT OR IGNORE INTO claims VALUES(?,?,?,?)',(scope,'socks_port',value,owner))
            c.execute('INSERT OR IGNORE INTO batches VALUES(?,?)', (plan['batch_id'], digest(plan)))

    def state(self, batch, key, phase):
        with self.db() as c:
            row = c.execute('SELECT status FROM phases WHERE batch=? AND node=? AND phase=?',
                            (batch, key, phase)).fetchone()
        return row[0] if row else 'pending'

    def reusable(self, batch, key, phase, dependency, request):
        if not dependency or self.state(batch, key, phase) != 'done':
            return False
        with self.db() as c:
            row = c.execute('SELECT dependency,expires,request FROM evidence WHERE batch=? AND node=? AND phase=?', (batch,key,phase)).fetchone()
        return bool(row and row[0] == dependency and time.time() <= row[1] <= time.time()+900 and row[2] == request)

    def proof(self, batch, key, phase, dependency, expires, request):
        with self.db() as c:
            c.execute('INSERT OR REPLACE INTO evidence VALUES(?,?,?,?,?,?)',(batch,key,phase,dependency,expires,request))

    def mark(self, batch, key, phase, status, attempt=False):
        previous=self.state(batch,key,phase)
        with self.db() as c:
            c.execute('''INSERT INTO phases VALUES(?,?,?,?,?,?) ON CONFLICT(batch,node,phase)
                DO UPDATE SET status=excluded.status,attempts=phases.attempts+?,updated=excluded.updated''',
                      (batch, key, phase, status, int(attempt), time.time(), int(attempt)))
        if previous!=status or attempt:
            fd=os.open(self.root/'events.jsonl',os.O_CREAT|os.O_WRONLY|os.O_APPEND,0o600)
            try:
                os.fchmod(fd,0o600)
                event=json.dumps({'batch':batch,'node':key,'phase':phase,'status':status,'time':time.time()})+'\n'
                os.write(fd,event.encode())
                if EVENTS:sys.stdout.write(event);sys.stdout.flush()
            finally:os.close(fd)


class Coordinator:
    def __init__(self, plan, root, jobs=4, timeout=900):
        self.plan = validate(plan)
        self.ledger = Ledger(root)
        self.jobs = jobs
        self.timeout = timeout
        self.dependencies = {}
        self.runtime = None
        self.ledger.reserve(plan)
        self.dir = self.ledger.root / plan['batch_id']
        self.dir.mkdir(mode=0o700, exist_ok=True)
        self.plan_file = self.dir / 'plan.json'
        atomic_json(self.plan_file, plan)
        if plan.get('runtime'):
            self.check_runtime()
            directory = str(Path(__file__).resolve().parent)
            if directory not in sys.path: sys.path.insert(0, directory)
            from site_adapter import SiteAdapter
            self.runtime = SiteAdapter(plan, self.dir, atomic_writer=atomic_json, resource_lock=lock)

    def check_runtime(self):
        runtime = self.plan['runtime']
        if runtime.get('kind') != 'ssh-site-v1': raise ValueError('Unknown runtime')
        if runtime.get('python') != str(Path(sys.executable).resolve()):raise ValueError('Use the frozen runtime interpreter')
        import paramiko
        if paramiko.__version__ != runtime.get('paramiko_version'):raise ValueError('Runtime dependency version changed')
        base = Path(__file__).resolve().parent
        required = [base/n for n in ('batch.py','site_adapter.py','ssh_pool.py','acceptance.py')]
        required += [Path(self.plan['context'][k]) for k in ('ssh_config','exit_registry')]
        files = runtime.get('files', {})
        if any(str(f) not in files for f in required): raise ValueError('Runtime source pin missing')
        for f, sha in files.items():
            if not Path(f).is_absolute() or hashlib.sha256(Path(f).read_bytes()).hexdigest() != sha:
                raise ValueError('Runtime code/config changed')


    def call(self, n, phase, action):
        fingerprint = digest({'context': self.plan['context'], 'node': n})
        path = self.dir / (n['key'] + '.json')
        atomic_json(path, dict(n, request_sha256=fingerprint))
        argv = self.plan['adapter'] + ['--plan', str(self.plan_file), '--node', str(path),
                                       '--phase', phase, '--action', action, '--state-dir', str(self.dir)]
        if self.runtime:
            self.check_runtime()
            r = self.runtime.call(dict(n, request_sha256=fingerprint), phase, action)
            if (r.get('key'),r.get('phase'),r.get('request_sha256')) != (n['key'],phase,fingerprint):
                raise RuntimeError('Runtime identity mismatch')
            if r.get('status') not in ('ready','absent','waiting','unknown','conflict'):
                raise RuntimeError('Runtime status invalid')
            return r
        with subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL,
                              start_new_session=True) as p:
            deadline = time.monotonic() + self.timeout
            while True:
                try:
                    output, _ = p.communicate(timeout=.2)
                    break
                except subprocess.TimeoutExpired:
                    if STOP.is_set() or time.monotonic() >= deadline:
                        os.killpg(p.pid, signal.SIGTERM)
                        try:
                            p.communicate(timeout=5)
                        except subprocess.TimeoutExpired:
                            os.killpg(p.pid, signal.SIGKILL)
                            p.communicate()
                        raise RuntimeError('Adapter interrupted; remote outcome must be probed')
        if p.returncode or len(output) > 65536:
            raise RuntimeError('Adapter failed; no stdout/stderr persisted')
        r = json.loads(output)
        if (r.get('key'), r.get('phase'), r.get('request_sha256')) != (n['key'], phase, fingerprint):
            raise RuntimeError('Adapter receipt identity mismatch')
        if r.get('status') not in ('ready', 'absent', 'waiting', 'unknown', 'conflict'):
            raise RuntimeError('Invalid adapter status')
        return r

    def phase(self, n, phase):
        try:
            return self._phase(n, phase)
        except Exception:
            self.ledger.mark(self.plan['batch_id'], n['key'], phase, 'unknown')
            return False

    def _phase(self, n, phase):
        batch, key = self.plan['batch_id'], n['key']
        resource = SHARED.get(phase, n['vps_id'])
        # Common panel/dns/relay lock names across batches and processes.
        namespace = self.plan['context']['panel_vps']
        if resource in ('dns', 'relay'):
            namespace = 'dpl-external'
        with lock(self.ledger.root, namespace + ':' + resource):
            dependency = self.dependencies.get(key, {}).get(phase)
            request = digest({'plan': self.plan, 'node': n})
            # Publish/audit/preflight are always live; unknown or running phases cannot use a cache.
            if phase in ('stage','dns','provision','relay','verify') and self.ledger.reusable(batch,key,phase,dependency,request):
                if self.runtime: self.check_runtime()
                return True
            r = self.call(n, phase, 'probe')
            if phase == 'preflight' and r['status'] == 'ready':
                deps = r.get('resume_dependencies', {})
                clean = {k:v for k,v in deps.items() if k in PHASES and isinstance(v,str) and re.fullmatch('[0-9a-f]{64}',v)}
                self.dependencies[key] = {k:digest({p:clean.get(p) for p in PHASES[:PHASES.index(k)+1]}) for k in clean}
            if r['status'] == 'ready':
                if dependency and type(r.get('valid_until')) in (int,float) and time.time() < r['valid_until'] <= time.time()+900:
                    self.ledger.proof(batch,key,phase,dependency,r['valid_until'],request)
                self.ledger.mark(batch, key, phase, 'done')
                return True
            if r['status'] == 'conflict':
                STOP.set()
            if phase in ('preflight', 'verify', 'audit') or r['status'] != 'absent' or r.get('retry_safe') is not True:
                self.ledger.mark(batch, key, phase, r['status'])
                return False
            if STOP.is_set():
                return False
            # Persist intent before a write; never infer success from the command exit code.
            self.ledger.mark(batch, key, phase, 'running', attempt=True)
            applied = self.call(n, phase, 'apply')
            if applied['status'] == 'conflict':
                self.ledger.mark(batch, key, phase, 'conflict')
                STOP.set()
                return False
            r = self.call(n, phase, 'probe')
            self.ledger.mark(batch, key, phase, 'done' if r['status'] == 'ready' else r['status'])
            if r['status'] == 'ready' and dependency and type(r.get('valid_until')) in (int,float) and time.time()<r['valid_until']<=time.time()+900:
                self.ledger.proof(batch,key,phase,dependency,r['valid_until'],request)
            if r['status'] == 'conflict':
                STOP.set()
            return r['status'] == 'ready'

    def worker(self, n):
        try:
            with lock(self.ledger.root, 'node:' + self.plan['context']['panel_vps'] + ':' + n['key']):
                self.ledger.mark(self.plan['batch_id'], n['key'], 'coordinator', 'running')
                # Already finished batches are re-audited; disabled QA need not be re-enabled.
                if self.ledger.state(self.plan['batch_id'], n['key'], 'audit') == 'done':
                    success = self.phase(n, 'audit')
                else:
                    success = True
                    for phase in PHASES:
                        if STOP.is_set() or not self.phase(n, phase):
                            success = False
                            break
                self.ledger.mark(self.plan['batch_id'], n['key'], 'coordinator', 'done' if success else 'waiting')
                return success
        except Exception:
            # Never print adapter payloads, commands, credentials or exception contents.
            self.ledger.mark(self.plan['batch_id'], n['key'], 'coordinator', 'unknown')
            return False

    def run(self):
        # Whole-batch barrier: no staging until every candidate passed current preflight.
        with concurrent.futures.ThreadPoolExecutor(max_workers=self.jobs) as pool:
            def check(n):
                try:
                    return self.phase(n, 'preflight')
                except Exception:
                    self.ledger.mark(self.plan['batch_id'], n['key'], 'preflight', 'unknown')
                    return False
            checks = list(pool.map(check, self.plan['nodes']))
            if not all(checks):
                return False
            return all(list(pool.map(self.worker, self.plan['nodes'])))


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('command', choices=('plan', 'run', 'status'))
    p.add_argument('--plan', required=True)
    p.add_argument('--state-root', required=True)
    p.add_argument('--jobs', type=int, default=4)
    p.add_argument('--timeout', type=int, default=900)
    p.add_argument('--events', action='store_true', help='Emit token-free phase transition JSON lines; final result is last line')
    p.add_argument('--execute', action='store_true', help='Required for a run that may write remotely')
    a = p.parse_args()
    global EVENTS
    EVENTS = a.events
    os.umask(0o077)
    if not 1 <= a.jobs <= 16 or not 1 <= a.timeout <= 3600:
        p.error('jobs must be 1..16; timeout 1..3600')
    try:
        plan = validate(read_json(a.plan))
        if a.command == 'plan':
            print(json.dumps({'batch_id': plan['batch_id'], 'nodes': len(plan['nodes']),
                              'plan_sha256': digest(plan), 'writes': False}))
            return 0
        if a.command == 'run' and not a.execute:
            p.error('run requires --execute; plan does not perform remote requests')
        if a.command == 'status':
            path = Path(a.state_root).resolve() / 'registry.sqlite3'
            rows = {}
            if path.exists():
                with sqlite3.connect(path.as_uri() + '?mode=ro', uri=True) as c:
                    prior = c.execute('SELECT digest FROM batches WHERE id=?', (plan['batch_id'],)).fetchone()
                    if prior and prior[0] != digest(plan):
                        raise ValueError('Batch plan differs from its reserved plan')
                    rows = {(key, phase): status for key, phase, status in c.execute(
                        'SELECT node,phase,status FROM phases WHERE batch=?', (plan['batch_id'],))}
            print(json.dumps({n['key']: {phase: rows.get((n['key'], phase), 'pending')
                                       for phase in ('preflight',) + PHASES + ('coordinator',)}
                              for n in plan['nodes']}))
            return 0
        runner = Coordinator(plan, a.state_root, a.jobs, a.timeout)
        for sig in (signal.SIGINT, signal.SIGTERM):
            signal.signal(sig, lambda *_: STOP.set())
        try:
            success = runner.run()
        finally:
            if runner.runtime:
                try:runner.runtime.close()
                except Exception:
                    for node in plan['nodes']:runner.ledger.mark(plan['batch_id'],node['key'],'coordinator','unknown')
                    raise
        print(json.dumps({'batch_id': plan['batch_id'], 'complete': success,
                          'next': None if success else 'Inspect status, resolve waiting/conflict, rerun the exact plan'}))
        return 0 if success else 2
    except Exception as exc:
        # Validation errors have no adapter/credential material.
        print(json.dumps({'error': str(exc) if isinstance(exc, ValueError) else 'Coordinator failed; inspect protected state'}))
        return 2


if __name__ == '__main__':
    sys.exit(main())