#!/usr/bin/env python3
"""
Twitch EventSub Webhook Bridge → Mock Discord Webhook
=====================================================
Runs via: python webhook_bridge.py
"""

import asyncio
import hmac
import hashlib
import json
import time
import random
import logging
from typing import Dict, Tuple

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
SECRET = "twitch_eventsub_shared_secret"
DISCORD_WEBHOOK_URL = "https://discord.com/api/webhooks/mock/mock"  # Mocked internally
IDEMPOTENCY_TTL = 300.0  # 5 minutes
RATE_LIMIT = 5.0         # forwards per second
RATE_CAPACITY = 5.0      # burst capacity

# ---------------------------------------------------------------------------
# Discord Mock
# ---------------------------------------------------------------------------
class DiscordRateLimitError(Exception):
    """Raised when the mock Discord endpoint returns a 429."""
    def __init__(self, retry_after: float):
        self.retry_after = retry_after
        super().__init__(f"Discord 429: Retry-After {retry_after:.2f}s")

async def discord_send(payload: dict) -> dict:
    """
    Mock Discord webhook endpoint.
    Succeeds ~95% of the time. Occasionally returns HTTP 429.
    """
    if random.random() < 0.05:  # 5% chance of rate limit
        retry_after = random.uniform(1.0, 3.0)
        raise DiscordRateLimitError(retry_after)
    await asyncio.sleep(0.05)  # Simulate network latency
    return {"status": 204}

# ---------------------------------------------------------------------------
# Idempotency Store
# ---------------------------------------------------------------------------
class IdempotencyStore:
    """In-memory deduplication store keyed by event_id."""
    def __init__(self, ttl: float = IDEMPOTENCY_TTL):
        self.store: Dict[str, float] = {}
        self.ttl = ttl

    def is_seen(self, event_id: str) -> bool:
        now = time.time()
        if event_id in self.store and now - self.store[event_id] < self.ttl:
            return True
        # Lazy eviction of stale entries
        self.store = {k: v for k, v in self.store.items() if now - v < self.ttl}
        return False

    def mark(self, event_id: str) -> None:
        self.store[event_id] = time.time()

# ---------------------------------------------------------------------------
# Token Bucket Rate Limiter
# ---------------------------------------------------------------------------
class TokenBucketLimiter:
    """Async token-bucket limiter. Caps rate to `rate` tokens/sec."""
    def __init__(self, rate: float, capacity: float):
        self.rate = rate
        self.capacity = capacity
        self.tokens = capacity
        self.last_refill = time.monotonic()
        self._lock = asyncio.Lock()

    async def acquire(self) -> None:
        while True:
            async with self._lock:
                now = time.monotonic()
                elapsed = now - self.last_refill
                self.tokens = min(self.capacity, self.tokens + elapsed * self.rate)
                self.last_refill = now

                if self.tokens >= 1.0:
                    self.tokens -= 1.0
                    return
                wait_time = (1.0 - self.tokens) / self.rate
            await asyncio.sleep(wait_time)

# ---------------------------------------------------------------------------
# Webhook Bridge Handler
# ---------------------------------------------------------------------------
class WebhookBridge:
    """Core bridge logic: HMAC verification, idempotency, rate limiting, forwarding."""
    def __init__(self):
        self.idempotency = IdempotencyStore()
        self.rate_limiter = TokenBucketLimiter(rate=RATE_LIMIT, capacity=RATE_CAPACITY)
        self.forward_count = 0
        self.forward_timestamps: list[float] = []

    async def handle_webhook(self, body: bytes) -> Tuple[int, str]:
        # 1. Parse JSON
        try:
            event = json.loads(body)
        except json.JSONDecodeError as e:
            return 400, f"Invalid JSON: {e}"

        event_id = event.get("event_id")
        if not event_id or not isinstance(event_id, str):
            return 400, "Missing or invalid event_id"

        # 2. Idempotency check
        if self.idempotency.is_seen(event_id):
            return 200, "Replayed"
        self.idempotency.mark(event_id)

        # 3. Rate limiting
        await self.rate_limiter.acquire()

        # 4. Forward to Discord with 429 backoff
        for attempt in range(2):  # 1 initial + 1 retry
            try:
                await discord_send({"event_id": event_id, "summary": "Forwarded"})
                self.forward_count += 1
                self.forward_timestamps.append(time.monotonic())
                return 200, "Forwarded"
            except DiscordRateLimitError as e:
                if attempt == 0:
                    logger.warning(f"Discord 429 detected. Backing off {e.retry_after:.2f}s...")
                    await asyncio.sleep(e.retry_after)
                else:
                    return 502, f"Discord 429 after retry: {e}"

        return 500, "Unexpected downstream error"

# ---------------------------------------------------------------------------
# Async HTTP Server
# ---------------------------------------------------------------------------
async def handle_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
    """Minimal async HTTP/1.1 handler for POST /webhook."""
    bridge = WEBHOOK_BRIDGE
    try:
        request_line = await reader.readline()
        if not request_line:
            return
        method, path, _ = request_line.decode().strip().split()
        if method != "POST" or path != "/webhook":
            resp = b"HTTP/1.1 404 Not Found\r\nContent-Length: 9\r\n\r\nNot Found"
            writer.write(resp)
            await writer.drain()
            return

        # Parse headers
        headers: Dict[str, str] = {}
        while True:
            line = await reader.readline()
            if not line or line == b"\r\n":
                break
            key, _, value = line.decode().strip().partition(":")
            headers[key.strip().lower()] = value.strip()

        # Parse body
        content_length = int(headers.get("content-length", 0))
        body = await reader.readexactly(content_length) if content_length > 0 else b""

        # HMAC Signature Verification
        provided_sig = headers.get("x-signature", "")
        if not provided_sig:
            status, message = 401, "Missing X-Signature header"
        else:
            computed_sig = hmac.new(SECRET.encode(), body, hashlib.sha256).hexdigest()
            if not hmac.compare_digest(computed_sig, provided_sig):
                status, message = 401, "Invalid signature"
            else:
                try:
                    status, message = await bridge.handle_webhook(body)
                except Exception as e:
                    logger.exception("Unhandled bridge error")
                    status, message = 500, f"Internal Error: {e}"

        # Send response
        response = (
            f"HTTP/1.1 {status} {message}\r\n"
            f"Content-Length: {len(message)}\r\n"
            f"Content-Type: text/plain\r\n"
            f"\r\n{message}"
        )
        writer.write(response.encode())
        await writer.drain()
    except asyncio.IncompleteReadError:
        logger.warning("Client disconnected prematurely")
    except Exception as e:
        logger.exception("Server connection error")
    finally:
        writer.close()
        try:
            await writer.wait_closed()
        except Exception:
            pass

# Global bridge instance shared across connections
WEBHOOK_BRIDGE = WebhookBridge()

# ---------------------------------------------------------------------------
# Test Client & Suite
# ---------------------------------------------------------------------------
async def send_webhook_request(event_id: str, signature: str = None, body_data: dict = None) -> Tuple[str, str]:
    """Helper to send a raw HTTP request and return (status_code, body)."""
    if signature is None:
        body_data = body_data or {"event_id": event_id, "type": "chat", "data": {}}
        raw_body = json.dumps(body_data).encode()
        signature = hmac.new(SECRET.encode(), raw_body, hashlib.sha256).hexdigest()
    else:
        raw_body = json.dumps(body_data or {}).encode()

    reader, writer = await asyncio.open_connection("127.0.0.1", 8080)
    request = (
        f"POST /webhook HTTP/1.1\r\n"
        f"Host: localhost\r\n"
        f"Content-Type: application/json\r\n"
        f"Content-Length: {len(raw_body)}\r\n"
        f"X-Signature: {signature}\r\n"
        f"\r\n"
    )
    writer.write(request.encode() + raw_body)
    await writer.drain()

    response = b""
    while b"\r\n\r\n" not in response:
        chunk = await reader.read(4096)
        if not chunk:
            break
        response += chunk
    header_end = response.find(b"\r\n\r\n")
    headers = response[:header_end].decode()
    body = response[header_end+4:].decode()
    status = headers.split("\r\n")[0].split()[1]
    writer.close()
    await writer.wait_closed()
    return status, body

async def run_tests():
    logger.info("Starting test suite...")
    server = await asyncio.start_server(handle_client, "127.0.0.1", 8080)
    await asyncio.sleep(0.2)  # Allow port binding

    try:
        # a) Correctly-signed request → forwarded once, returns 200
        status, body = await send_webhook_request("test-a")
        logger.info(f"a) Valid signature: {status} {body}")
        assert status == "200" and body == "Forwarded", f"Failed a: {status} {body}"
        assert WEBHOOK_BRIDGE.forward_count == 1, f"Expected 1 forward, got {WEBHOOK_BRIDGE.forward_count}"

        # b) Tampered signature → 401, nothing forwarded
        status, body = await send_webhook_request("test-b", signature="invalidsig")
        logger.info(f"b) Tampered signature: {status} {body}")
        assert status == "401", f"Failed b: {status}"
        assert WEBHOOK_BRIDGE.forward_count == 1, "Forward count should not increase on 401"

        # c) Replay same event_id within 5 min → skipped, returns 200
        status, body = await send_webhook_request("test-a")
        logger.info(f"c) Replay event: {status} {body}")
        assert status == "200" and body == "Replayed", f"Failed c: {status} {body}"
        assert WEBHOOK_BRIDGE.forward_count == 1, "Forward count should not increase on replay"

        # d) Burst >5 events in one second → rate limiter delays excess
        WEBHOOK_BRIDGE.forward_count = 0
        WEBHOOK_BRIDGE.forward_timestamps.clear()
        burst_ids = [f"burst-{i}" for i in range(10)]
        start = time.monotonic()
        tasks = [send_webhook_request(rid) for rid in burst_ids]
        results = await asyncio.gather(*tasks)
        elapsed = time.monotonic() - start

        all_200 = all(r[0] == "200" for r in results)
        logger.info(f"d) Burst 10 events: all 200? {all_200}, elapsed {elapsed:.2f}s, forwards {WEBHOOK_BRIDGE.forward_count}")
        assert all_200, "Burst requests failed"
        assert WEBHOOK_BRIDGE.forward_count == 10, f"Expected 10 forwards, got {WEBHOOK_BRIDGE.forward_count}"
        assert elapsed > 1.0, f"Rate limiter failed: took {elapsed:.2f}s, expected >1s"

        logger.info("\n✅ All tests passed!")
    except AssertionError as e:
        logger.error(f"\n❌ Test failed: {e}")
    except Exception as e:
        logger.exception(f"\n❌ Unexpected error: {e}")
    finally:
        server.close()
        await server.wait_closed()

if __name__ == "__main__":
    asyncio.run(run_tests())