#!/usr/bin/env python3
"""Bounded CPU-quota experiment. Only transient user services, loopback HTTP."""
import argparse
import concurrent.futures as cf
import hashlib
import http.client
import json
import math
import os
from pathlib import Path
import platform
import random
import socket
import statistics
import subprocess
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

PAYLOAD = b'kamui-cpu-quota-lab-' * 48
PORT = 18767

def work(n):
    x = PAYLOAD
    for _ in range(n):
        x = hashlib.sha256(x).digest()
    return x

def emit(value):
    print(json.dumps(value, ensure_ascii=False), flush=True)

def cgroup_path():
    return Path('/sys/fs/cgroup' + Path('/proc/self/cgroup').read_text().strip().split('::')[1])

def values(path):
    return {k: int(v) for k, v in (line.split() for line in path.read_text().splitlines())}

def host_sample():
    cpu = [int(x) for x in Path('/proc/stat').read_text().splitlines()[0].split()[1:9]]
    mem = {line.split(':')[0]: int(line.split()[1]) for line in Path('/proc/meminfo').read_text().splitlines()}
    thermal = Path('/sys/class/thermal/thermal_zone0/temp')
    return {'at': time.time(), 'cpu_total': sum(cpu), 'cpu_idle': cpu[3] + cpu[4],
            'available_mib': mem['MemAvailable'] / 1024,
            'temp_c': int(thermal.read_text()) / 1000 if thermal.exists() else None,
            'load': list(os.getloadavg())}

def guard(stop, samples, failures):
    previous = host_sample()
    while not stop.wait(1):
        now = host_sample()
        elapsed = now['cpu_total'] - previous['cpu_total']
        now['cpu_percent'] = 100 * (1 - (now['cpu_idle'] - previous['cpu_idle']) / elapsed) if elapsed else 0
        samples.append(now)
        previous = now
        reason = None
        if now['available_mib'] < 1536:
            reason = 'available memory below 1536 MiB'
        if now['cpu_percent'] > 50:
            reason = 'host CPU above 50%'
        if now['temp_c'] is not None and now['temp_c'] >= 72:
            reason = 'temperature reached 72 C'
        if reason:
            failures.append(reason)
            stop.set()
            return

def serve(iterations):
    class Handler(BaseHTTPRequestHandler):
        protocol_version = 'HTTP/1.1'
        def setup(self):
            super().setup()
            self.connection.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        def log_message(self, *args):
            pass
        def do_GET(self):
            start = time.thread_time_ns()
            if self.path == '/work':
                work(iterations)
            elif self.path != '/health':
                self.send_error(404)
                return
            body = json.dumps({'cpu_ns': time.thread_time_ns() - start}).encode()
            self.send_response(200)
            self.send_header('Content-Length', str(len(body)))
            self.send_header('Content-Type', 'application/json')
            self.end_headers()
            self.wfile.write(body)
    server = ThreadingHTTPServer(('127.0.0.1', PORT), Handler)
    server.daemon_threads = True
    server.request_queue_size = 16
    cg = cgroup_path()
    emit({'ready': True, 'cgroup': str(cg), 'cpu_max': (cg / 'cpu.max').read_text().strip(),
          'memory_max': (cg / 'memory.max').read_text().strip(), 'pid': os.getpid()})
    server.serve_forever(poll_interval=0.2)

def percentile(items, p):
    a = sorted(items)
    return a[max(0, math.ceil(p * len(a)) - 1)]

def measure(args, quota, pattern, repeat, iterations, seq, stop):
    unit = f'kamui-cpu-lab-{os.getpid()}-{seq}'
    command = ['systemd-run', '--user', '--wait', '--pipe', '--collect', '--quiet',
               '--unit=' + unit, '--property=CPUQuota=' + str(quota * 100) + '%',
               '--property=CPUQuotaPeriodSec=100ms', '--property=MemoryMax=128M',
               '--property=TasksMax=32', '--property=RuntimeMaxSec=35',
               '--property=NoNewPrivileges=yes',
               sys.executable, str(Path(__file__).resolve()), '--server', '--iterations', str(iterations)]
    child = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    local = threading.local()
    connections = []
    lock = threading.Lock()
    def request(due, record=True):
        if stop.is_set():
            raise RuntimeError('safety guard stopped experiment')
        if not hasattr(local, 'connection'):
            local.connection = http.client.HTTPConnection('127.0.0.1', PORT, timeout=2)
            with lock:
                connections.append(local.connection)
        begin = time.perf_counter()
        local.connection.request('GET', '/work')
        response = local.connection.getresponse()
        body = json.loads(response.read())
        end = time.perf_counter()
        if response.status != 200:
            raise RuntimeError('unexpected HTTP status')
        return {'scheduled_s': due, 'start_s': begin, 'end_s': end,
                'latency_ms': (end - due) * 1000, 'http_ms': (end - begin) * 1000,
                'start_lag_ms': (begin - due) * 1000, 'work_cpu_ms': body['cpu_ns'] / 1e6}
    try:
        line = child.stdout.readline()
        if not line:
            raise RuntimeError('server startup failed: ' + child.stderr.read())
        info = json.loads(line)
        if info['cpu_max'] != f'{int(quota * 100000)} 100000' or info['memory_max'] != '134217728':
            raise RuntimeError('resource caps did not match requested settings')
        cg = Path(info['cgroup'])
        with cf.ThreadPoolExecutor(max_workers=8) as pool:
            warmup = [pool.submit(request, time.perf_counter(), False) for _ in range(8)]
            for f in warmup:
                f.result(timeout=3)
            time.sleep(0.4 + random.Random(seq).uniform(0, 0.1))
            before = values(cg / 'cpu.stat')
            stats_start = time.perf_counter()
            host_before = host_sample()
            start = time.perf_counter() + 0.1
            futures = []
            for i in range(round(args.duration * 20)):
                offset = i / 20 if pattern == 'uniform' else (i // 8) * 0.4
                due = start + offset
                delay = due - time.perf_counter()
                if delay > 0 and stop.wait(delay):
                    raise RuntimeError('safety guard stopped experiment')
                if stop.is_set() or time.perf_counter() - due > 0.25:
                    raise RuntimeError('safety stop or generator fell behind')
                if sum(not f.done() for f in futures[-24:]) >= 16:
                    raise RuntimeError('more than 16 pending requests')
                futures.append(pool.submit(request, due))
            samples = [f.result(timeout=3) for f in futures]
            end = start + args.duration
            if time.perf_counter() < end:
                stop.wait(end - time.perf_counter())
            if stop.is_set():
                raise RuntimeError('safety guard stopped experiment')
            after = values(cg / 'cpu.stat')
            elapsed = time.perf_counter() - start
            host_after = host_sample()
            latency = [r['latency_ms'] for r in samples]
            delta = {k: after[k] - before[k] for k in before}
            result = {'quota_cores': quota, 'pattern': pattern, 'repeat': repeat,
                      'duration_s': elapsed, 'expected': round(args.duration * 20), 'successful': len(samples),
                      'rps': len(samples) / elapsed, 'latency_p50_ms': percentile(latency, .50),
                      'latency_p95_ms': percentile(latency, .95), 'latency_p99_ms': percentile(latency, .99),
                      'latency_max_ms': max(latency), 'generator_lag_p99_ms': percentile([r['start_lag_ms'] for r in samples], .99),
                      'work_cpu_p50_ms': statistics.median(r['work_cpu_ms'] for r in samples),
                      'cpu_cores': delta['usage_usec'] / ((time.perf_counter() - stats_start) * 1e6), 'cpu_stat_delta': delta,
                      'server_info': info, 'host_before': host_before, 'host_after': host_after,
                      'memory_peak_bytes': int((cg / 'memory.peak').read_text()) if (cg / 'memory.peak').exists() else None,
                      'memory_events': values(cg / 'memory.events'), 'requests': samples}
            return result
    finally:
        for connection in connections:
            connection.close()
        # Stops only this experiment's transient unit, never an existing service.
        subprocess.run(['systemctl', '--user', 'stop', unit], capture_output=True, timeout=8)
        try:
            child.communicate(timeout=5)
        except subprocess.TimeoutExpired:
            child.terminate()
            child.communicate(timeout=3)

def main(args):
    if not 1 <= args.duration <= 24 or not 1 <= args.repeats <= 3:
        raise SystemExit('duration/repeats outside safe bounds')
    # Refuse to share the chosen loopback port with an existing application.
    probe = socket.socket()
    probe.bind(('127.0.0.1', PORT))
    probe.close()
    base = host_sample()
    if base['temp_c'] is not None and base['temp_c'] >= 72 or base['available_mib'] < 1536:
        raise SystemExit('unsafe baseline')
    stop, host_samples, failures = threading.Event(), [], []
    watcher = threading.Thread(target=guard, args=(stop, host_samples, failures), daemon=True)
    watcher.start()
    timer = threading.Timer(360, lambda: (failures.append('overall 360-second watchdog'), stop.set()))
    timer.start()
    result = {'metadata': {'platform': platform.platform(), 'python': platform.python_version(),
                          'cpus': os.cpu_count(), 'baseline': base, 'rate': 20, 'burst_size': 8,
                          'period_us': 100000, 'duration': args.duration, 'repeats': args.repeats,
                          'safety': {'max_cores': .5, 'memory_mib': 128, 'threads': 8, 'pending': 16,
                                     'max_temp_c': 72, 'max_host_cpu_percent': 50, 'min_available_mib': 1536}},
              'runs': [], 'host_samples': host_samples, 'failures': failures}
    try:
        n = 1000
        for _ in range(3):
            t = time.thread_time()
            work(n)
            n = max(1, round(n * .004 / (time.thread_time() - t)))
        calibration = []
        for _ in range(7):
            t = time.thread_time()
            work(n)
            calibration.append((time.thread_time() - t) * 1000)
        result['metadata'].update(iterations=n, calibration_cpu_ms=calibration)
        emit({'event': 'calibrated', 'iterations': n, 'cpu_ms': statistics.median(calibration)})
        cells = [(.25, 'uniform'), (.25, 'burst'), (.5, 'uniform'), (.5, 'burst')]
        seq = 0
        for repeat in range(1, args.repeats + 1):
            # Reverse then rotate ordering to reduce a fixed warmup/temperature ordering bias.
            order = cells if repeat == 1 else list(reversed(cells)) if repeat == 2 else cells[2:] + cells[:2]
            for quota, pattern in order:
                if stop.is_set():
                    raise RuntimeError('safety guard stopped experiment')
                seq += 1
                run = measure(args, quota, pattern, repeat, n, seq, stop)
                result['runs'].append(run)
                Path(args.output).write_text(json.dumps(result, ensure_ascii=False, indent=2))
                emit({k: v for k, v in run.items() if k not in ('requests', 'server_info', 'host_before', 'host_after', 'memory_events')})
                if stop.wait(2):
                    raise RuntimeError('safety guard stopped experiment')
    except Exception as exc:
        failures.append(str(exc))
        raise
    finally:
        timer.cancel()
        stop.set()
        watcher.join(timeout=2)
        Path(args.output).write_text(json.dumps(result, ensure_ascii=False, indent=2))
        emit({'event': 'finished', 'runs': len(result['runs']), 'failures': failures,
              'max_temp_c': max([s['temp_c'] for s in host_samples if s['temp_c'] is not None], default=None),
              'max_host_cpu_percent': max([s['cpu_percent'] for s in host_samples], default=None)})

if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('--server', action='store_true')
    p.add_argument('--iterations', type=int, default=1000)
    p.add_argument('--duration', type=int, default=20)
    p.add_argument('--repeats', type=int, default=3)
    p.add_argument('--output', default='results.json')
    a = p.parse_args()
    serve(a.iterations) if a.server else main(a)