#!/usr/bin/env python3
"""TCP Stream black-box test: use against a running RustProxy and mock backends.

The 'selftest' mode validates this harness only; it does NOT run RustProxy.
All traffic uses loopback by default. Do not point it at a production service.
"""
from __future__ import annotations

import argparse
from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed
import secrets
import socket
import socketserver
import sys
import threading
import time

MAX_LINE = 1024


class TCPServer(socketserver.ThreadingTCPServer):
    allow_reuse_address = True
    daemon_threads = True
    request_queue_size = 128


class TCPBackend(socketserver.StreamRequestHandler):
    def handle(self) -> None:
        self.request.settimeout(10)
        backend_id = self.server.backend_id.encode('ascii')
        try:
            while True:
                line = self.rfile.readline(MAX_LINE + 1)
                if not line or len(line) > MAX_LINE or not line.endswith(b'\n'):
                    return
                self.wfile.write(backend_id + b'|' + line)
                self.wfile.flush()
        except (OSError, TimeoutError):
            return


def run_server(host: str, port: int, backend_id: str) -> None:
    with TCPServer((host, port), TCPBackend) as server:
        server.backend_id = backend_id
        print(f'TCP mock {backend_id} listening on {server.server_address}', flush=True)
        try:
            server.serve_forever(poll_interval=0.1)
        except KeyboardInterrupt:
            pass


def check_one(index: int, host: str, port: int, messages: int, timeout: float, hold_ms: float, nonce: str) -> str:
    backend = None
    with socket.create_connection((host, port), timeout=timeout) as connection:
        connection.settimeout(timeout)
        with connection.makefile('rb') as reader:
            for seq in range(messages):
                payload = f'rp-tcp:{nonce}:{index}:{seq}\n'.encode('ascii')
                connection.sendall(payload)
                response = reader.readline(MAX_LINE + 1)
                try:
                    observed, echoed = response.split(b'|', 1)
                except ValueError as exc:
                    raise AssertionError(f'session {index}: missing mock backend response: {response!r}') from exc
                assert echoed == payload, f'session {index}: response mismatch: {response!r}'
                assert observed, f'session {index}: empty backend ID'
                if backend is None:
                    backend = observed
                assert observed == backend, f'session {index}: TCP backend changed during connection'
                if hold_ms and seq != messages - 1:
                    time.sleep(hold_ms / 1000)
    return backend.decode('ascii')


def run_client(host: str, port: int, sessions: int, messages: int, parallel: int, timeout: float,
               hold_ms: float, min_backends: int) -> Counter:
    nonce = secrets.token_hex(8)
    started = time.monotonic()
    counts: Counter[str] = Counter()
    with ThreadPoolExecutor(max_workers=parallel) as pool:
        futures = [pool.submit(check_one, i, host, port, messages, timeout, hold_ms, nonce)
                   for i in range(sessions)]
        for future in as_completed(futures):
            counts[future.result()] += 1
    assert len(counts) >= min_backends, f'observed {len(counts)} backend(s), required {min_backends}: {dict(counts)}'
    print(f'TCP PASS: {sessions} sessions, {sessions * messages} messages, '
          f'{len(counts)} backend(s), {time.monotonic() - started:.2f}s; assignments={dict(sorted(counts.items()))}')
    return counts


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest='mode', required=True)
    server = sub.add_parser('server', help='run a local TCP mock backend')
    server.add_argument('--host', default='127.0.0.1')
    server.add_argument('--port', type=int, required=True)
    server.add_argument('--backend-id', required=True, help='unique ASCII marker (e.g. node-a)')
    client = sub.add_parser('client', help='test forwarding via an already-running RustProxy listener')
    client.add_argument('--host', default='127.0.0.1', help='RustProxy TCP listener host')
    client.add_argument('--port', type=int, required=True, help='RustProxy TCP listener port')
    client.add_argument('--sessions', type=int, default=24)
    client.add_argument('--messages', type=int, default=3)
    client.add_argument('--parallel', type=int, default=8)
    client.add_argument('--timeout', type=float, default=4.0)
    client.add_argument('--hold-ms', type=float, default=20.0, help='pause between messages to keep TCP sessions alive')
    client.add_argument('--min-backends', type=int, default=1, help='minimum distinct backend IDs observed')
    sub.add_parser('selftest', help='test mock peer + harness directly, WITHOUT RustProxy')
    args = parser.parse_args()
    if args.mode == 'server':
        if not args.backend_id.isascii() or not args.backend_id or '|' in args.backend_id or '\n' in args.backend_id:
            parser.error('--backend-id must be nonempty ASCII without | or newline')
        run_server(args.host, args.port, args.backend_id)
    elif args.mode == 'client':
        if min(args.port, args.sessions, args.messages, args.parallel, args.min_backends) < 1 or args.port > 65535 or args.timeout <= 0 or args.hold_ms < 0:
            parser.error('invalid port, count, timeout or hold delay')
        run_client(args.host, args.port, args.sessions, args.messages, args.parallel,
                   args.timeout, args.hold_ms, args.min_backends)
    else:
        with TCPServer(('127.0.0.1', 0), TCPBackend) as server:
            server.backend_id = 'local-mock'
            thread = threading.Thread(target=server.serve_forever, daemon=True)
            thread.start()
            try:
                run_client('127.0.0.1', server.server_address[1], 8, 3, 4, 2.0, 5, 1)
                print('SELFTEST PASS: tested mock TCP peer, not RustProxy')
            finally:
                server.shutdown()
                thread.join(timeout=2)


if __name__ == '__main__':
    try:
        main()
    except (OSError, AssertionError, UnicodeError) as exc:
        print(f'TCP FAIL: {exc}', file=sys.stderr)
        sys.exit(1)