#!/usr/bin/env python3
"""Disposable, loopback-only observability lab. Standard library only."""
import argparse
import json
import secrets
import threading
import time
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import urlparse

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--port', type=int, default=8090)
parser.add_argument('--worker-delay', type=float, default=0.2)
parser.add_argument('--otlp', default='http://127.0.0.1:4318/v1/traces')
parser.add_argument('--break-context', action='store_true')
parser.add_argument('--fail-jobs', action='store_true')
args = parser.parse_args()
if not 0 <= args.worker_delay <= 10:
    parser.error('worker-delay must be between 0 and 10 seconds')
lock = threading.Lock()
jobs = {}
stats = {'accepted': 0, 'success': 0, 'failed': 0, 'seconds': 0.0, 'export_failed': 0}

def log(event, **fields):
    print(json.dumps({'time_ns': time.time_ns(), 'event': event, 'service': 'dispatch-lab', **fields}), flush=True)

def span(name, trace_id, span_id, start, end, parent=None):
    item = {'traceId': trace_id, 'spanId': span_id, 'name': name,
            'kind': 1, 'startTimeUnixNano': str(start), 'endTimeUnixNano': str(end)}
    if parent:
        item['parentSpanId'] = parent
    payload = {'resourceSpans': [{'resource': {'attributes': [
        {'key': 'service.name', 'value': {'stringValue': 'dispatch-lab'}}]},
        'scopeSpans': [{'scope': {'name': 'dispatch-learning'}, 'spans': [item]}]}]}
    try:
        request = urllib.request.Request(args.otlp, data=json.dumps(payload).encode(),
            headers={'Content-Type': 'application/json'}, method='POST')
        with urllib.request.urlopen(request, timeout=2) as response:
            response.read()
        log('telemetry_exported', trace_id=trace_id, span_id=span_id)
    except Exception as exc:
        with lock:
            stats['export_failed'] += 1
        log('telemetry_export_failed', error=type(exc).__name__, trace_id=trace_id)

def worker(job_id, trace_id, parent_id):
    if args.break_context:
        trace_id, parent_id = secrets.token_hex(16), None
    start, child_id = time.time_ns(), secrets.token_hex(8)
    log('worker_started', job_id=job_id, trace_id=trace_id)
    time.sleep(args.worker_delay)
    outcome = 'failed' if args.fail_jobs else 'success'
    end = time.time_ns()
    with lock:
        jobs[job_id]['state'] = outcome
        jobs[job_id]['worker_trace_id'] = trace_id
        stats[outcome] += 1
        stats['seconds'] += (end-start)/1e9
    log('worker_finished', job_id=job_id, trace_id=trace_id, outcome=outcome)
    span('build export', trace_id, child_id, start, end, parent_id)

class Handler(BaseHTTPRequestHandler):
    def respond(self, status, value, content_type='application/json'):
        body = value.encode() if isinstance(value, str) else json.dumps(value).encode()
        self.send_response(status)
        self.send_header('Content-Type', content_type)
        self.send_header('Content-Length', str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_POST(self):
        if urlparse(self.path).path != '/exports':
            return self.respond(404, {'error': 'unknown route'})
        start = time.time_ns()
        trace_id, parent_id, job_id = secrets.token_hex(16), secrets.token_hex(8), secrets.token_hex(4)
        with lock:
            jobs[job_id] = {'state': 'pending', 'accepted_trace_id': trace_id}
            stats['accepted'] += 1
        log('export_accepted', job_id=job_id, trace_id=trace_id)
        self.respond(202, {'job_id': job_id, 'trace_id': trace_id, 'status_url': '/exports/'+job_id})
        end = time.time_ns()
        threading.Thread(target=worker, args=(job_id, trace_id, parent_id), daemon=True).start()
        span('POST /exports', trace_id, parent_id, start, end)

    def do_GET(self):
        path = urlparse(self.path).path
        if path == '/health':
            return self.respond(200, {'status': 'listening'})
        if path == '/metrics':
            with lock:
                s = dict(stats)
                pending = sum(j['state'] == 'pending' for j in jobs.values())
            text = ('# TYPE dispatch_exports_accepted_total counter\n'
                    f'dispatch_exports_accepted_total {s["accepted"]}\n'
                    '# TYPE dispatch_exports_completed_total counter\n'
                    f'dispatch_exports_completed_total{{outcome="success"}} {s["success"]}\n'
                    f'dispatch_exports_completed_total{{outcome="failed"}} {s["failed"]}\n'
                    '# TYPE dispatch_exports_pending gauge\n'
                    f'dispatch_exports_pending {pending}\n'
                    '# TYPE dispatch_worker_duration_seconds_total counter\n'
                    f'dispatch_worker_duration_seconds_total {s["seconds"]}\n'
                    '# TYPE dispatch_telemetry_export_failures_total counter\n'
                    f'dispatch_telemetry_export_failures_total {s["export_failed"]}\n')
            return self.respond(200, text, 'text/plain; version=0.0.4')
        if path.startswith('/exports/'):
            with lock:
                job = dict(jobs.get(path.rsplit('/',1)[-1], {}))
            return self.respond(200 if job else 404, job or {'error': 'unknown export'})
        self.respond(404, {'error': 'unknown route'})

    def log_message(self, *_):
        pass

log('service_started', port=args.port, worker_delay=args.worker_delay)
try:
    ThreadingHTTPServer(('127.0.0.1', args.port), Handler).serve_forever()
except KeyboardInterrupt:
    log('service_stopped')