# webhook_bridge.py
# Python 3.11+  stdlib only
# Run: python webhook_bridge.py

import os, sys, json, hmac, hashlib, time, threading, queue, asyncio, random
import http.server, http.client

SECRET = os.getenv('WEBHOOK_SECRET', 'supersecret')

# ---------- Idempotency ----------
event_store = {}
store_lock = threading.Lock()

def evict_stale():
    now = time.monotonic()
    with store_lock:
        stale = [k for k, v in event_store.items() if now - v > 300]
        for k in stale:
            del event_store[k]

# ---------- Rate limiter: token bucket 5/s ----------
CAPACITY = 5
RATE = 5.0
tokens = CAPACITY
last_refill = time.monotonic()
bucket_lock = threading.Lock()

def acquire_token():
    global tokens, last_refill
    while True:
        with bucket_lock:
            now = time.monotonic()
            elapsed = now - last_refill
            tokens = min(CAPACITY, tokens + elapsed * RATE)
            last_refill = now
            if tokens >= 1:
                tokens -= 1
                return
            wait = (1 - tokens) / RATE
        time.sleep(wait)

# ---------- Mock Discord ----------
async def discord_send(payload):
    await asyncio.sleep(0.005)
    if random.random() < 0.05:          # ~5% 429
        class Resp:
            status = 429
            headers = {'Retry-After': '1'}
        return Resp()
    class Resp:
        status = 200
        headers = {}
    return Resp()

# ---------- Forwarding worker ----------
forward_queue = queue.Queue()
forward_log = []          # {event_id, ts}
log_lock = threading.Lock()

def worker():
    while True:
        item = forward_queue.get()
        if item is None:
            break
        event_id, payload = item
        try:
            acquire_token()
            resp = asyncio.run(discord_send(payload))
            if resp.status == 429:
                retry = int(resp.headers.get('Retry-After', '1'))
                print(f"[WARN] 429 for {event_id}, retry after {retry}s", file=sys.stderr)
                time.sleep(retry)
                resp = asyncio.run(discord_send(payload))
            if resp.status == 200:
                with log_lock:
                    forward_log.append({'event_id': event_id, 'ts': time.monotonic()})
            else:
                print(f"[ERROR] Discord forward failed {event_id} status {resp.status}", file=sys.stderr)
        except Exception as e:
            print(f"[ERROR] Discord forward exception {e}", file=sys.stderr)
        finally:
            forward_queue.task_done()

# ---------- HTTP handler ----------
class Handler(http.server.BaseHTTPRequestHandler):
    def do_POST(self):
        if self.path != '/webhook':
            self.send_response(404); self.end_headers(); return

        length = int(self.headers.get('Content-Length', 0))
        body = self.rfile.read(length)

        sig = self.headers.get('X-Signature')
        if not sig:
            self.send_response(401); self.end_headers(); self.wfile.write(b'Missing signature'); return

        expected = hmac.new(SECRET.encode(), body, hashlib.sha256).hexdigest()
        if not hmac.compare_digest(expected, sig):
            self.send_response(401); self.end_headers(); self.wfile.write(b'Invalid signature'); return

        try:
            data = json.loads(body)
        except json.JSONDecodeError:
            self.send_response(400); self.end_headers(); self.wfile.write(b'Invalid JSON'); return

        event_id = data.get('event_id')
        event_type = data.get('type')
        if not event_id or not event_type:
            self.send_response(400); self.end_headers(); self.wfile.write(b'Missing fields'); return

        now = time.monotonic()
        with store_lock:
            evict_stale()
            if event_id in event_store and now - event_store[event_id] < 300:
                self.send_response(200); self.end_headers(); self.wfile.write(b'OK replay'); return
            event_store[event_id] = now

        payload = {'content': f'Event {event_type} received', 'event_id': event_id}
        forward_queue.put((event_id, payload))
        self.send_response(200); self.end_headers(); self.wfile.write(b'OK')

    def log_message(self, fmt, *args):
        print(f"{self.client_address[0]} - {fmt%args}")

# ---------- Tests ----------
def make_request(event_id, event_type, secret=SECRET, tamper=False):
    body = json.dumps({'event_id': event_id, 'type': event_type, 'data': {}}).encode()
    sig = hmac.new(secret.encode(), body, hashlib.sha256).hexdigest()
    if tamper:
        sig = '0'*64
    conn = http.client.HTTPConnection('localhost', 8000, timeout=5)
    conn.request('POST', '/webhook', body, {'Content-Type':'application/json','X-Signature':sig})
    resp = conn.getresponse()
    resp.read(); conn.close()
    return resp.status

def reset_state():
    with store_lock:
        event_store.clear()
    with log_lock:
        forward_log.clear()

def run_tests():
    time.sleep(0.2)
    print('Test a) correct signature')
    reset_state()
    s = make_request('id1','chat')
    time.sleep(0.1)
    assert s == 200, f'expected 200 got {s}'
    assert len(forward_log) == 1, 'should be forwarded once'
    print('  ok')

    print('Test b) tampered signature')
    s = make_request('id2','chat', tamper=True)
    assert s == 401, f'expected 401 got {s}'
    assert len(forward_log) == 1, 'should not forward'
    print('  ok')

    print('Test c) replay idempotency')
    s = make_request('id1','chat')
    assert s == 200
    time.sleep(0.05)
    assert len(forward_log) == 1, 'replay must not increase count'
    print('  ok')

    print('Test d) rate limiter burst')
    reset_state()
    for i in range(10):
        make_request(f'burst{i}','chat')
    # wait for worker to drain
    time.sleep(2.5)
    assert len(forward_log) == 10, 'all 10 should eventually forward'
    ts = [e['ts'] for e in forward_log]
    elapsed = ts[-1] - ts[0]
    assert elapsed >= 1.5, f'rate limiter not enforced, elapsed {elapsed:.2f}s'
    print(f'  ok, elapsed {elapsed:.2f}s >=1.5s')
    print('All tests passed.')

if __name__ == '__main__':
    worker_thread = threading.Thread(target=worker, daemon=True)
    worker_thread.start()

    server = http.server.ThreadingHTTPServer(('localhost', 8000), Handler)
    srv_thread = threading.Thread(target=server.serve_forever, daemon=True)
    srv_thread.start()

    try:
        run_tests()
    finally:
        server.shutdown()