import concurrent.futures as cf
import datetime as dt
import http.server
import importlib.metadata
import json
import os
from pathlib import Path
import random
import threading
import time

import httpx
import modal
from e2b_code_interpreter import Sandbox as E2B
from daytona import Daytona, DaytonaConfig, CreateSandboxFromSnapshotParams

ROOT = Path('/app/results')
ROOT.mkdir(exist_ok=True)
LOCK = threading.Lock()
STATE = {'status': 'initializing', 'round': 0}
N = int(os.environ.get('BENCH_RUNS', '300'))
START_ROUND = 0
MARKER = 'arena-startup-ready-20261001'
COMMAND = 'python3 -c "print(\'arena-startup-ready-20261001\')"'
PROVIDERS = ['lizard', 'e2b', 'daytona', 'modal']
CLIENTS = {}

def emit(rec):
    rec['recorded_at'] = dt.datetime.now(dt.timezone.utc).isoformat()
    line = json.dumps(rec)
    with LOCK:
        with (ROOT / 'raw.jsonl').open('a') as f:
            f.write(line + '\n')
        print(line, flush=True)

def clean_error(e):
    value = f'{type(e).__name__}: {e}'
    if isinstance(e, httpx.HTTPStatusError):
        value += ' response=' + e.response.text[:1000]
    for k, v in os.environ.items():
        if any(x in k for x in ('TOKEN', 'KEY', 'SECRET', 'PASSWORD')) and len(v) > 5:
            value = value.replace(v, '[redacted]')
    return value[:1500]

def prepare():
    CLIENTS['lizard'] = httpx.Client(
        base_url=os.environ.get('LIZARD_API_URL', 'https://lizard.build').rstrip('/'),
        headers={'X-API-Key': os.environ['LIZARD_API_KEY']}, timeout=90)
    CLIENTS['daytona'] = Daytona(DaytonaConfig(api_key=os.environ['DAYTONA_API_KEY'], target='us'))
    CLIENTS['modal_app'] = modal.App.lookup('sandbox-arena', create_if_missing=True)
    CLIENTS['modal_image'] = modal.Image.debian_slim(python_version='3.12')

def create(p):
    if p == 'lizard':
        r = CLIENTS[p].post('/api/sandboxes', json={
            'template': 'base', 'size': 'small', 'projectId': os.environ['LIZARD_PROJECT_ID'],
            'region': 'us-east-1', 'timeoutMs': 300000})
        r.raise_for_status()
        return r.json()
    if p == 'e2b':
        return E2B.create(api_key=os.environ['E2B_API_KEY'], timeout=300,
                          metadata={'benchmark': 'arena-startup-20261012-small-python'}, request_timeout=90)
    if p == 'daytona':
        return CLIENTS[p].create(CreateSandboxFromSnapshotParams(
            language='python', labels={'benchmark': 'arena-startup-20261012-small-python'},
            auto_stop_interval=5, auto_delete_interval=1), timeout=90)
    return modal.Sandbox.create(app=CLIENTS['modal_app'], image=CLIENTS['modal_image'],
        cpu=(2, 2), memory=(2048, 2048), region='us-east', timeout=300)

def sid(p, sb):
    if p == 'lizard': return sb['sandboxId']
    if p == 'e2b': return sb.sandbox_id
    if p == 'daytona': return sb.id
    return sb.object_id

def execute(p, sb, cmd):
    if p == 'lizard':
        r = CLIENTS[p].post('/api/sandboxes/' + sid(p, sb) + '/exec',
                           json={'cmd': cmd, 'timeoutMs': 30000}, timeout=45)
        r.raise_for_status()
        d = r.json()
        return d['exitCode'], d.get('stdout', ''), d.get('stderr', '')
    if p == 'e2b':
        r = sb.commands.run(cmd, timeout=30)
        return r.exit_code, r.stdout, r.stderr
    if p == 'daytona':
        r = sb.process.exec(cmd, timeout=30)
        return r.exit_code, r.result or '', ''
    r = sb.exec('sh', '-c', cmd, timeout=30)
    out, err = r.stdout.read(), r.stderr.read()
    r.wait()
    return r.returncode, out, err

def destroy(p, sb):
    if p == 'lizard':
        r = CLIENTS[p].delete('/api/sandboxes/' + sid(p, sb))
        if r.status_code != 404: r.raise_for_status()
    elif p == 'e2b': sb.kill()
    elif p == 'daytona': CLIENTS[p].delete(sb, timeout=60)
    else: sb.terminate()

def trial(p, index, phase):
    rec = {'provider': p, 'sample': index, 'phase': phase, 'ok': False,
           'started_at': dt.datetime.now(dt.timezone.utc).isoformat()}
    sb = None
    stage = 'create'
    start = time.perf_counter()
    try:
        sb = create(p)
        created = time.perf_counter()
        rec.update(sandbox_id=sid(p, sb), create_ms=(created-start)*1000)
        # Persist the ID immediately, so interrupted runs can clean up their resources.
        emit({'phase': 'allocated', 'provider': p, 'sample': index,
              'sandbox_id': sid(p, sb), 'create_ms': rec['create_ms']})
        stage = 'first_exec'
        code, stdout, stderr = execute(p, sb, COMMAND)
        ready = time.perf_counter()
        rec.update(first_exec_ms=(ready-created)*1000, ready_ms=(ready-start)*1000,
                   exit_code=code)
        if code != 0 or stdout.strip() != MARKER:
            raise RuntimeError(f'Invalid readiness response: exit={code}, stdout={stdout[:200]!r}, stderr={stderr[:200]!r}')
        rec['ok'] = True
        if phase == 'warmup' and index == 0:
            inspect = "python3 --version; cat /sys/fs/cgroup/cpu.max /sys/fs/cgroup/memory.max 2>/dev/null; cat /sys/fs/cgroup/cpu/cpu.cfs_quota_us /sys/fs/cgroup/cpu/cpu.cfs_period_us /sys/fs/cgroup/memory/memory.limit_in_bytes 2>/dev/null; true"
            rec['environment'] = execute(p, sb, inspect)[1]
            if p == 'lizard':
                rec['allocation'] = {k: sb.get(k) for k in ('cpus', 'memoryMb', 'region', 'template', 'size')}
            if p == 'daytona':
                rec['allocation'] = {k: str(getattr(sb, k, None)) for k in ('cpu', 'memory', 'target', 'snapshot')}
    except Exception as e:
        rec.update(error=clean_error(e), error_stage=stage, elapsed_ms=(time.perf_counter()-start)*1000)
    finally:
        if sb is not None:
            for attempt in range(3):
                try:
                    t = time.perf_counter()
                    destroy(p, sb)
                    rec['cleanup_ms'] = (time.perf_counter()-t)*1000
                    rec['cleanup_ok'] = True
                    break
                except Exception as e:
                    rec['cleanup_error'] = clean_error(e)
                    rec['cleanup_ok'] = False
                    time.sleep(1 + attempt)
        rec['finished_at'] = dt.datetime.now(dt.timezone.utc).isoformat()
        emit(rec)
    return rec

def run():
    try:
        prepare()
        emit({'phase': 'metadata', 'n': N, 'command': COMMAND, 'providers': PROVIDERS,
              'runner_region': os.environ.get('RAILWAY_REPLICA_REGION'),
              'versions': {k: importlib.metadata.version(k) for k in
                           ['httpx','modal','e2b','e2b-code-interpreter','daytona']},
              'method': 'one active sandbox per provider; 3 warmups; seeded shuffled rounds; no measurement retries; prepared images; perf_counter'})
        rng = random.Random(20261001)
        with cf.ThreadPoolExecutor(max_workers=4) as pool:
            for phase, count in [('warmup', 3), ('measured', N)]:
                STATE['status'] = phase
                consecutive = {p: 0 for p in PROVIDERS}
                for i in range(START_ROUND, count):
                    STATE['round'] = i + 1
                    order = PROVIDERS.copy()
                    rng.shuffle(order)
                    futures = [pool.submit(trial, p, i, phase) for p in order]
                    results = [f.result() for f in futures]
                    for rec in results:
                        p = rec['provider']
                        consecutive[p] = 0 if rec['ok'] else consecutive[p] + 1
                        if rec.get('cleanup_ok') is False:
                            raise RuntimeError(f'Cleanup failed for {p}; stopping new allocations')
                    if phase == 'warmup' and any(r.get('error_stage') == 'first_exec' and not r['ok'] for r in results):
                        raise RuntimeError('Warmup readiness failed; inspect template command support before measuring')
                    # Capacity refusals are measured outcomes; at one create every round
                    # they do not justify silently truncating the requested sample size.
                    if any('401 Unauthorized' in r.get('error', '') or '403 Forbidden' in r.get('error', '') for r in results):
                        raise RuntimeError('Authentication/authorization failed; stopping new allocations')
                    time.sleep(.25)
        STATE['status'] = 'done'
    except Exception as e:
        STATE.update(status='failed', error=clean_error(e))
    emit({'phase': 'finished', **STATE})

class Handler(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        # No public domain is configured. SSH can read the raw file directly.
        data = json.dumps(STATE).encode()
        self.send_response(200)
        self.end_headers()
        self.wfile.write(data)
    def log_message(self, *args): pass

if __name__ == '__main__':
    threading.Thread(target=run, daemon=True).start()
    http.server.HTTPServer(('0.0.0.0', int(os.environ.get('PORT', '8080'))), Handler).serve_forever()
