#!/usr/bin/env python3
"""Founder commands for one bounded-execution run (NK-WP-0058-T04). Each prints one value-free line.

  BX="python3 tools/bounded-execution-run/bx.py"
  $BX ready admissions/BEA-YYYY-NNNN.yaml   # readiness; one platform-admin login (OpenBao status)
  $BX request                              # file the approval; the id goes into the state file
  (approve the memo in informed-decision: browser profile of admin, aal2 with OTP)
  $BX run                                  # preflight, then window open, run, close; one platform-admin login
  $BX status                               # the run's state and receipts, no login
  $BX close                                # emergency: delete the write role; safe to repeat

Inputs live in a state file (~/.local/state/net-kingdom/bounded-execution/<id>.json), not in shell variables.
The admission named last is the current one; --admission selects another. Receipts are numbered by the tool in
docs/evidence/<date>-<id>/. Each command is safe to repeat: `run` never starts a second run for an admission.
The OpenBao steps run as silent children of `warden access openbao-platform-admin-login --exec`; the window
is closed in the same child, also when the run fails or is interrupted.

--attest names the trust states you attest for the run (default: runtime_secret_trust,runtime_identity_trust).
They are recorded in the state file and the run config, never edited in by hand.
"""
import argparse
import datetime as dt
import hashlib
import json
import os
import re
import subprocess
import sys
import time
from pathlib import Path

import yaml

REPO = Path(__file__).resolve().parents[2]
HOME = Path.home()
STATE_DIR = Path(os.environ.get('BX_STATE_DIR', HOME / '.local/state/net-kingdom/bounded-execution'))
RAPP = Path(os.environ.get('BX_RAPP', REPO.parent / 'rapp-bounded-executor'))
PLATFORM = Path(os.environ.get('BX_PLATFORM', REPO.parent / 'railiance-platform'))
KUBECONFIG = Path(os.environ.get('BX_KUBECONFIG', HOME / '.kube/config-railiance01'))
BAO_ADDR = os.environ.get('BAO_ADDR', 'http://127.0.0.1:18200')
LANE = ['warden', 'access', 'openbao-platform-admin-login', '--exec', '--']
LANES_TOOL = REPO / 'tools/service-binds/executor-lanes.py'
READY_TOOL = REPO / 'tools/bounded-execution-readiness/ready.py'
CCRS = {'write': 'CCR-2026-0058-*.yaml', 'journey': 'CCR-2026-0059-*.yaml'}
DELIVERIES = ('keycape-nk-bind-executor-approval-client', 'bounded-executor-tokens')
ATTEST = 'runtime_secret_trust,runtime_identity_trust'
WARDEN_EXIT = {10: 'login_failed', 11: 'child_failed', 12: 'child_output', 13: 'revoke_unconfirmed',
               14: 'cleanup_failed'}


class Runner:
    """Every outside command goes through here; tests replace it."""

    def run(self, argv, env=None, timeout=900):
        full = {**os.environ, **(env or {})}
        done = subprocess.run([str(a) for a in argv], capture_output=True, env=full, timeout=timeout)
        return done.returncode, done.stdout.decode(errors='replace')

    def sleep(self, seconds):
        time.sleep(seconds)


def kubectl():
    return f'kubectl --kubeconfig {KUBECONFIG}'


def json_blocks(text):
    """The value-free JSON results in a log: one-line objects and the executor's indented result blocks."""
    out = []
    for block in re.findall(r'^\{\n.*?\n\}$', text, re.M | re.S) + re.findall(r'^\{.*\}$', text, re.M):
        try:
            out.append(json.loads(block))
        except ValueError:
            pass
    return out


class Bx:
    def __init__(self, runner=None, now=None):
        self.r = runner or Runner()
        self.now = now or (lambda: dt.datetime.now(dt.timezone.utc))

    # ------------------------------------------------------------------ state and receipts
    def state_path(self, admission_id):
        return STATE_DIR / f'{admission_id}.json'

    def load(self, admission=None):
        if admission:
            doc = yaml.safe_load(Path(admission).read_text())
            sid = doc['spec']['id']
            path = self.state_path(sid)
            if path.exists():
                return json.loads(path.read_text())
            evidence = REPO / 'docs/evidence' / f"{self.now():%Y-%m-%d}-{sid.lower()}"
            return {'admission_id': sid, 'admission': str(Path(admission).resolve()), 'evidence': str(evidence),
                    'approval_id': None, 'attest': None, 'run': None}
        current = STATE_DIR / 'current'
        if not current.exists():
            raise SystemExit(self.line({'status': 'refused', 'reason': 'no_admission', 'next': 'bx ready ADMISSION'}))
        return json.loads(self.state_path(current.read_text().strip()).read_text())

    def save(self, state):
        STATE_DIR.mkdir(parents=True, exist_ok=True)
        path = self.state_path(state['admission_id'])
        path.write_text(json.dumps(state, indent=2, sort_keys=True) + '\n')
        (STATE_DIR / 'current').write_text(state['admission_id'] + '\n')

    def receipt(self, state, name):
        folder = Path(state['evidence'])
        folder.mkdir(parents=True, exist_ok=True)
        numbers = [int(p.name[:2]) for p in folder.glob('[0-9][0-9]-*')]
        return folder / f'{max(numbers, default=0) + 1:02d}-{name}.json'

    def write(self, path, data):
        path.write_text(json.dumps(data, indent=2, sort_keys=True) + '\n')
        return path

    def config(self, state):
        """The run config: the rapp's pilot config with the attested trust states. Non-secret."""
        cfg = json.loads((RAPP / 'examples/config.pilot.json').read_text())
        cfg['trust_satisfied'] = list(state['attest'])
        path = STATE_DIR / f"{state['admission_id']}.config.json"
        path.write_text(json.dumps(cfg, indent=2) + '\n')
        return path

    @staticmethod
    def line(data):
        return json.dumps(data, separators=(', ', ': '))

    def lane(self, child):
        """Run a silent child under one platform-admin login. The child gets an isolated HOME from warden, so every
        path it needs is passed as an absolute path here."""
        env = {'BAO_ADDR': BAO_ADDR, 'VAULT_ADDR': BAO_ADDR, 'PYTHONWARNINGS': 'ignore', 'BX_STATE_DIR': STATE_DIR,
               'BX_RAPP': RAPP, 'BX_PLATFORM': PLATFORM, 'BX_KUBECONFIG': KUBECONFIG, 'NK_PLATFORM_ROOT': PLATFORM}
        code, _ = self.r.run(LANE + child, env={k: str(v) for k, v in env.items()})
        return code

    # ------------------------------------------------------------------ commands
    def ready(self, admission, attest):
        state = self.load(admission)
        state['attest'] = state['attest'] or attest.split(',')
        self.save(state)
        rec = self.receipt(state, 'readiness')
        self.r.run(['python3', READY_TOOL, '--admission', state['admission'],
                    '--approval-id', state['approval_id'] or 'none', '--receipt', rec])
        agent = (json.loads(rec.read_text()) if rec.exists()
                 else {'agent_checks_pass': False, 'failures': ['no_receipt']})
        custody = {}
        for name in DELIVERIES:
            code, out = self.r.run(kubectl().split() + ['-n', 'sso', 'get', 'externalsecret', name, '-o',
                                   'jsonpath={.status.conditions[?(@.type=="Ready")].status}'])
            custody[name] = code == 0 and out.strip() == 'True'
        self.write(self.receipt(state, 'custody'), {'action': 'custody-delivery', 'ready': custody,
                                                     'values_recorded': False})
        status = self.receipt(state, 'lanes-status')
        warden = self.lane(['python3', LANES_TOOL, 'status', '--receipt', status, '--quiet'])
        lanes = json.loads(status.read_text()) if status.exists() else {}
        failures = list(agent.get('failures', [])) + [f'custody:{n}' for n, ok in custody.items() if not ok]
        if not lanes.get('ready_for_window'):
            failures.append('openbao:' + (lanes.get('reason') or WARDEN_EXIT.get(warden, 'not_ready_for_window')))
        state['ready'] = {'at': f'{self.now():%Y-%m-%dT%H:%M:%SZ}', 'failures': failures}
        self.save(state)
        nxt = 'bx run' if state['approval_id'] else 'bx request'
        return self.line({'ready': not failures, 'admission': state['admission_id'], 'failures': failures,
                          'next': nxt if not failures else 'fix the failures, then bx ready again'})

    def request(self):
        state = self.load()
        if state['approval_id']:
            return self.line({'approval_id': state['approval_id'], 'status': 'already_requested',
                              'next': 'approve as admin (aal2), then bx run'})
        code, out = self.r.run(['make', '-C', RAPP, 'request', f"ADMISSION={state['admission']}",
                                f'CONFIG={self.config(state)}', f'KUBECTL={kubectl()}'])
        result = next((b for b in reversed(json_blocks(out)) if 'outcome' in b), {})
        self.write(self.receipt(state, 'request'), {**result, 'exit': code})
        if result.get('approval_id'):
            state['approval_id'] = result['approval_id']
            self.save(state)
        return self.line({'approval_id': result.get('approval_id'), 'outcome': result.get('outcome'),
                          'next': 'open the memo as admin (aal2, fresh OTP), approve, then bx run'
                          if result.get('outcome') == 'presented' else 'bx request again (idempotent)'})

    def run(self):
        state = self.load()
        if state.get('run') and state['run'].get('started'):
            return self.line({**self.summary(state), 'status': 'already_run'})
        if not state['approval_id']:
            return self.line({'status': 'refused', 'reason': 'no_approval', 'next': 'bx request'})
        self.config(state)
        code = self.lane(['python3', Path(__file__).resolve(), '_run-inner', '--state',
                          self.state_path(state['admission_id'])])
        state = json.loads(self.state_path(state['admission_id']).read_text())
        if code and not (state.get('run') or {}).get('finished'):
            state['run'] = {**(state.get('run') or {}), 'warden': WARDEN_EXIT.get(code, f'exit_{code}')}
            self.save(state)
        return self.line(self.summary(state))

    def run_inner(self, state_file):
        """The silent lane child: preflight, then window open, run, poll, and close in `finally`. Prints nothing."""
        state = json.loads(Path(state_file).read_text())
        cfg = STATE_DIR / f"{state['admission_id']}.config.json"
        k = kubectl()
        run = state['run'] = {'started': False, 'receipts': []}

        def note(path):
            run['receipts'].append(Path(path).name)
            self.save(state)

        _, out = self.r.run(['make', '-C', RAPP, 'preflight', f"ADMISSION={state['admission']}",
                             f"APPROVAL_ID={state['approval_id']}", f'CONFIG={cfg}', f'KUBECTL={k}'])
        report = next((b for b in reversed(json_blocks(out)) if b.get('mode') == 'preflight-only'),
                      {'ready': False, 'failures': ['preflight_no_report']})
        note(self.write(self.receipt(state, 'preflight'), report))
        if not report.get('ready'):
            run.update(outcome='not_started', reason='preflight_not_ready', failures=report.get('failures'))
            return self.save(state)
        digests = {}
        for lane, pattern in CCRS.items():
            files = sorted((PLATFORM / 'credential-change-requests').glob(pattern))
            digests[lane] = hashlib.sha256(files[0].read_bytes()).hexdigest()
        opened = self.receipt(state, 'window-open')
        try:
            self.r.run(['python3', LANES_TOOL, 'window-open', '--admission', state['admission'],
                        '--write-ccr-sha256', digests['write'], '--journey-ccr-sha256', digests['journey'],
                        '--receipt', opened, '--quiet'])
            note(opened)
            if json.loads(opened.read_text()).get('status') != 'window_opened':
                run.update(outcome='not_started', reason='window_not_opened')
                return
            run['started'] = True
            self.save(state)
            code, _ = self.r.run(['make', '-C', RAPP, 'run', f"ADMISSION={state['admission']}",
                                  f"APPROVAL_ID={state['approval_id']}", f'CONFIG={cfg}', f'KUBECTL={k}'])
            if code:
                run.update(outcome='not_started', reason='job_not_created')
                return
            job = 'bounded-executor-' + state['admission_id'].lower()
            budget = yaml.safe_load(Path(state['admission']).read_text())['spec']['budget']['wall_time_seconds']
            for _ in range(int(budget / 5) + 60):
                _, st = self.r.run(k.split() + ['-n', 'sso', 'get', f'job/{job}', '-o',
                                                'jsonpath={.status.succeeded}{.status.failed}'])
                if st.strip():
                    break
                self.r.sleep(5)
            _, logs = self.r.run(k.split() + ['-n', 'sso', 'logs', f'job/{job}'])
            result = next((b for b in reversed(json_blocks(logs)) if 'outcome' in b),
                          {'outcome': 'unknown', 'reason': 'no_result_in_log'})
            note(self.write(self.receipt(state, 'result'), result))
            run.update(outcome=result.get('outcome'), reason=result.get('reason'), run_id=result.get('run_id'))
        finally:
            closed = self.receipt(state, 'window-close')
            self.r.run(['python3', LANES_TOOL, 'window-close', '--admission', state['admission'],
                        '--receipt', closed, '--quiet'])
            if closed.exists():
                note(closed)
                run['window_closed'] = json.loads(closed.read_text()).get('status') == 'window_closed'
            run['finished'] = True
            self.save(state)

    def summary(self, state):
        run = state.get('run') or {}
        return {'admission': state['admission_id'], 'approval_id': state['approval_id'],
                'outcome': run.get('outcome'), 'reason': run.get('reason'), 'window_closed': run.get('window_closed'),
                'warden': run.get('warden'), 'error': run.get('error'), 'receipts': len(run.get('receipts', [])),
                'evidence': os.path.relpath(state['evidence'], REPO)}

    def status(self):
        return self.line(self.summary(self.load()))

    def close(self):
        state = self.load()
        rec = self.receipt(state, 'window-close')
        code = self.lane(['python3', LANES_TOOL, 'window-close', '--admission', state['admission'], '--receipt', rec,
                          '--quiet'])
        closed = json.loads(rec.read_text()) if rec.exists() else {}
        return self.line({'status': closed.get('status', 'failed'), 'reason': closed.get('reason'),
                          'warden': WARDEN_EXIT.get(code) if code else None})


def main(argv=None, bx=None):
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument('command', choices=['ready', 'request', 'run', 'status', 'close', '_run-inner'])
    p.add_argument('admission', nargs='?')
    p.add_argument('--attest', default=ATTEST)
    p.add_argument('--state', help=argparse.SUPPRESS)
    a = p.parse_args(argv)
    bx = bx or Bx()
    if a.command == '_run-inner':  # a silent lane child: any output would fail the lane, so record, never print
        try:
            bx.run_inner(a.state)
        except Exception as error:
            state = json.loads(Path(a.state).read_text())
            state['run'] = {**(state.get('run') or {}), 'error': type(error).__name__, 'finished': True}
            Path(a.state).write_text(json.dumps(state, indent=2, sort_keys=True) + '\n')
        return 0
    if a.command == 'ready':
        if not a.admission:
            p.error('ready takes the admission file')
        print(bx.ready(a.admission, a.attest))
        return 0
    if a.admission:
        bx.save(bx.load(a.admission))
    print({'request': bx.request, 'run': bx.run, 'status': bx.status, 'close': bx.close}[a.command]())
    return 0


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