#!/usr/bin/env python3
"""Collect only selected status fields over the existing railiance01 SSH lane."""
import argparse
from datetime import datetime, timezone
import json
import re
import os
import tempfile
from pathlib import Path
import shlex
import subprocess
import sys

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / 'scripts'))
from service_assurance import read_contract, admission


def query(args, allowed_codes=(0,)):
    result = subprocess.run(['ssh', '-o', 'BatchMode=yes', '-o', 'ConnectTimeout=10',
        'railiance01', shlex.join(['kubectl', '--request-timeout=10s', *args])],
        capture_output=True, text=True, timeout=25)
    if result.returncode not in allowed_codes:
        raise ValueError('query unavailable')
    return json.loads(result.stdout)


def memory_bytes(value):
    match = re.fullmatch(r"([0-9]+(?:\.[0-9]+)?)(Ki|Mi|Gi|Ti|K|M|G|T)?", value)
    if not match:
        raise ValueError('invalid memory quantity')
    unit = match.group(2) or ''
    multipliers = {'': 1, 'Ki': 1024, 'Mi': 1024**2, 'Gi': 1024**3,
        'Ti': 1024**4, 'K': 1000, 'M': 1000**2, 'G': 1000**3, 'T': 1000**4}
    return float(match.group(1)) * multipliers[unit]


def capture():
    contract = read_contract(ROOT / 'assurance/service-contract.json')
    uid = query(['get', 'namespace', 'kube-system', '-o', 'go-template={{printf "%q" .metadata.uid}}'])
    if uid != contract['cluster_uid']:
        raise ValueError('wrong cluster')
    signals = {}
    def add(name, result, when=None):
        signals[name] = {'result': result, 'observed_at': when or datetime.now(timezone.utc).isoformat()}
    baseline = json.loads((ROOT / 'assurance/admission-baseline.json').read_text())
    if admission() != baseline:
        raise ValueError('admission baseline drift')
    for cell in ('apps-pg', 'platform-pg', 'platform-pg-2'):
        names = [cell + '.' + suffix for suffix in ('ready', 'backup', 'wal')]
        try:
            d = query(['get', 'cluster', cell, '-n', 'databases', '-o',
                'go-template={"ready":{{.status.readyInstances}},"instances":{{.spec.instances}},"lastBackup":{{printf "%q" .status.lastSuccessfulBackup}},"conditions":[{{range $i,$v := .status.conditions}}{{if $i}},{{end}}{"type":{{printf "%q" $v.type}},"status":{{printf "%q" $v.status}}}{{end}}]}'])
            add(names[0], 'pass' if d['ready'] == d['instances'] and d['ready'] > 0 else 'fail')
            if d['lastBackup']:
                add(names[1], 'pass', d['lastBackup'])
            add(names[2], 'pass' if any(c == {'type': 'ContinuousArchiving', 'status': 'True'} for c in d['conditions']) else 'fail')
        except (ValueError, KeyError, TypeError, subprocess.TimeoutExpired):
            for name in names:
                add(name, 'unavailable')
        try:
            metrics = query(['get', '--raw', '/apis/metrics.k8s.io/v1beta1/namespaces/databases/pods/' + cell + '-1'])
            if not metrics['containers']:
                raise ValueError('empty metrics')
            memory = sum(memory_bytes(c['usage']['memory']) for c in metrics['containers'])
            # Same read-only pg_stat_activity aggregate as the package ops surface,
            # narrowed to a count: no SQL text, usernames, database rows or values.
            count = query(['exec', '-n', 'databases', cell + '-1', '-c', 'postgres', '--',
                'psql', '-U', 'postgres', '-d', 'postgres', '-Atqc',
                'SELECT count(*) FROM pg_stat_activity'])
            limits = baseline['cells'][cell]
            add(cell + '.headroom', 'pass' if type(count) is int and count >= 0
                and count <= limits['max_connections'] * 0.8
                and memory <= memory_bytes(limits['memory_limit']) * 0.8 else 'fail',
                metrics['timestamp'])
        except (OSError, ValueError, KeyError, TypeError, subprocess.TimeoutExpired):
            add(cell + '.headroom', 'unavailable')
    try:
        status = query(['exec', '-n', 'openbao', 'openbao-0', '--',
            'bao', 'status', '-format=json'], allowed_codes=(0, 2))
        add('openbao.seal', 'pass' if status.get('sealed') is False else 'fail')
    except (OSError, ValueError, KeyError, TypeError, subprocess.TimeoutExpired):
        add('openbao.seal', 'unavailable')
    try:
        rows = query(['get', 'externalsecrets', '-A', '-o',
            'go-template=[{{range $i,$v := .items}}{{if $i}},{{end}}{"refresh":{{printf "%q" $v.status.refreshTime}},"conditions":[{{range $j,$c := $v.status.conditions}}{{if $j}},{{end}}{"type":{{printf "%q" $c.type}},"status":{{printf "%q" $c.status}}}{{end}}]}{{end}}]'])
        ready = bool(rows) and all(any(c == {'type': 'Ready', 'status': 'True'} for c in row['conditions']) for row in rows)
        # Empty/missing refresh is a failure, never an apparently healthy zero age.
        refresh = min((r['refresh'] for r in rows), default='')
        add('eso.ready', 'pass' if ready else 'fail')
        if refresh:
            add('eso.refresh', 'pass', refresh)
    except (ValueError, KeyError, TypeError, subprocess.TimeoutExpired):
        add('eso.ready', 'unavailable')
        add('eso.refresh', 'unavailable')
    # No token, Secret, application data/logs or seal/unseal mutation.
    # Native restore and offsite receipts remain separate attended evidence.
    return {'schema': 'railiance-platform.observation.v1', 'cluster_uid': uid,
        'captured_at': datetime.now(timezone.utc).isoformat(), 'signals': signals}


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', type=Path)
    args = parser.parse_args()
    try:
        data = json.dumps(capture(), indent=2) + '\n'
        if args.output:
            temporary = None
            try:
                with tempfile.NamedTemporaryFile(mode='w', dir=args.output.parent, delete=False) as f:
                    temporary = f.name
                    f.write(data)
                    f.flush()
                    os.fsync(f.fileno())
                os.replace(temporary, args.output)
                temporary = None
            finally:
                if temporary:
                    os.unlink(temporary)
            print(json.dumps({'captured': True}))
        else:
            print(data, end='')
    except (OSError, ValueError, KeyError, TypeError, subprocess.TimeoutExpired):
        print(json.dumps({'schema': 'railiance-platform.assurance-error.v1', 'error': 'capture-unavailable-or-wrong-cluster'}))
        sys.exit(2)