"""backnforth server — static site + WebSocket relay, Python stdlib only.

    python3 server.py [--no-browser]

Serves the therapist page on http://<ip>:8000/ and the session relay on
ws://<ip>:8787. No pip installs, no virtualenv — any Python 3.7+ works,
including old machines that will never see a package index.
"""
import base64
import hashlib
import json
import os
import socket
import socketserver
import struct
import sys
import threading
import time
import webbrowser
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer

HTTP_PORT = 8000
WS_PORT = 8787
_WS_GUID = b'258EAFA5-E914-47DA-95CA-C5AB0DC85B11'

SESSIONS = {}   # session id -> set of live WSHandler
LOCK = threading.Lock()
_SERVERS = []
TOKENS = None   # None = open relay; a set = hosts must present a valid access code
STATS = None    # path to a JSONL file; None = no metrics at all


def anon(ip):
    """Daily-salted hash: counts uniques per day without ever storing an IP."""
    day = time.strftime('%Y-%m-%d')
    return hashlib.sha1((day + '|' + (ip or '?')).encode()).hexdigest()[:10]


def log_stat(event, **fields):
    if not STATS:
        return
    fields.update(t=int(time.time()), e=event)
    try:
        with open(STATS, 'a', encoding='utf-8') as f:
            f.write(json.dumps(fields) + '\n')
    except OSError:
        pass


def local_ip():
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(('8.8.8.8', 80))
        ip = s.getsockname()[0]
        s.close()
        return ip
    except OSError:
        try:
            return socket.gethostbyname(socket.gethostname())
        except OSError:
            return '127.0.0.1'


def shutdown_all():
    for srv in _SERVERS:
        try:
            srv.shutdown()
        except Exception:
            pass


class WSHandler(socketserver.BaseRequestHandler):
    """Minimal RFC6455 server: unfragmented text frames, ping/pong, close.

    ponytail: no fragmentation/binary support — the state payloads are a few
    hundred bytes of JSON and browsers never fragment at that size.
    """

    def setup(self):
        self.session_id = None
        self.send_lock = threading.Lock()
        self.hdrs = {}
        self.role = None
        self.joined_at = None

    def _client_ip(self):
        xff = self.hdrs.get(b'x-forwarded-for')
        if xff:
            return xff.decode('utf-8', 'replace').split(',')[0].strip()
        return self.client_address[0]

    def handle(self):
        if not self._handshake():
            return
        try:
            while True:
                msg = self._read_message()
                if msg is None:
                    break
                self._dispatch(msg)
        except (ConnectionError, OSError):
            pass

    def finish(self):
        if self.session_id:
            with LOCK:
                SESSIONS.get(self.session_id, set()).discard(self)
            if self.joined_at:
                log_stat('session', role=self.role or '?',
                         secs=int(time.time() - self.joined_at))

    # ---- relay protocol (same JSON messages as the old websockets server) ----

    def _dispatch(self, text):
        try:
            data = json.loads(text)
        except ValueError:
            return
        t = data.get('type')
        if t == 'join':
            sid = data.get('sessionId') or 'default'
            if TOKENS is not None:
                # hosts need an access code; patients may only join a LIVE session
                # (they got the random room id from their therapist's link)
                if data.get('role') == 'patient':
                    with LOCK:
                        live = bool(SESSIONS.get(sid))
                    if not live:
                        self._send_json({'type': 'join_denied', 'reason': 'session not live'})
                        return
                elif data.get('token') not in TOKENS:
                    self._send_json({'type': 'join_denied', 'reason': 'access code required'})
                    return
            self.session_id = sid
            self.role = data.get('role') or '?'
            self.joined_at = time.time()
            with LOCK:
                SESSIONS.setdefault(self.session_id, set()).add(self)
            self._send_json({'type': 'joined', 'sessionId': self.session_id})
            log_stat('join', role=self.role, u=anon(self._client_ip()))
        elif t in ('state', 'patientViewport'):
            if TOKENS is not None and not self.session_id:
                return  # hardened relay only forwards for joined connections
            sid = (self.session_id if TOKENS is not None
                   else data.get('sessionId') or self.session_id or 'default')
            out = json.dumps({'type': t, 'payload': data.get('payload')})
            with LOCK:
                peers = list(SESSIONS.get(sid, ()))
            for peer in peers:
                if peer is not self:
                    try:
                        peer._send_text(out)
                    except Exception:
                        pass
        elif t in ('whoami', 'detect_ip'):
            self._send_json({'type': 'detect_ip', 'ip': local_ip()})
        elif t == 'ping':
            self._send_json({'type': 'pong'})
        elif t == 'shutdown':
            if self.client_address[0] in ('127.0.0.1', '::1', local_ip()):
                self._send_json({'type': 'shutdown_ack'})
                threading.Thread(target=shutdown_all, daemon=True).start()

    # ---- websocket plumbing ----

    def _handshake(self):
        self.request.settimeout(10)
        data = b''
        while b'\r\n\r\n' not in data:
            chunk = self.request.recv(4096)
            if not chunk or len(data) > 65536:
                return False
            data += chunk
        headers = {}
        for line in data.split(b'\r\n')[1:]:
            if b':' in line:
                k, v = line.split(b':', 1)
                headers[k.strip().lower()] = v.strip()
        self.hdrs = headers
        key = headers.get(b'sec-websocket-key')
        if not key:
            return False
        accept = base64.b64encode(hashlib.sha1(key + _WS_GUID).digest()).decode()
        self.request.sendall((
            'HTTP/1.1 101 Switching Protocols\r\n'
            'Upgrade: websocket\r\n'
            'Connection: Upgrade\r\n'
            'Sec-WebSocket-Accept: ' + accept + '\r\n\r\n'
        ).encode())
        self.request.settimeout(None)
        return True

    def _recv_exact(self, n):
        buf = b''
        while len(buf) < n:
            chunk = self.request.recv(n - len(buf))
            if not chunk:
                return None
            buf += chunk
        return buf

    def _read_message(self):
        while True:
            head = self._recv_exact(2)
            if head is None:
                return None
            b1, b2 = head
            opcode = b1 & 0x0F
            length = b2 & 0x7F
            if length == 126:
                ext = self._recv_exact(2)
                if ext is None:
                    return None
                length = struct.unpack('>H', ext)[0]
            elif length == 127:
                ext = self._recv_exact(8)
                if ext is None:
                    return None
                length = struct.unpack('>Q', ext)[0]
            if length > 1_000_000:
                return None
            mask = self._recv_exact(4) if (b2 & 0x80) else b'\x00\x00\x00\x00'
            if mask is None:
                return None
            payload = self._recv_exact(length) if length else b''
            if payload is None:
                return None
            payload = bytes(c ^ mask[i % 4] for i, c in enumerate(payload))
            if opcode == 0x1:    # text
                try:
                    return payload.decode('utf-8')
                except UnicodeDecodeError:
                    return None
            elif opcode == 0x9:  # ping -> pong
                self._send_frame(0xA, payload)
            elif opcode == 0x8:  # close
                try:
                    self._send_frame(0x8, b'')
                except Exception:
                    pass
                return None
            # anything else (binary/continuation): skip

    def _send_frame(self, opcode, payload):
        n = len(payload)
        head = bytes([0x80 | opcode])
        if n < 126:
            head += bytes([n])
        elif n < 65536:
            head += bytes([126]) + struct.pack('>H', n)
        else:
            head += bytes([127]) + struct.pack('>Q', n)
        with self.send_lock:
            self.request.sendall(head + payload)

    def _send_text(self, text):
        self._send_frame(0x1, text.encode('utf-8'))

    def _send_json(self, obj):
        self._send_text(json.dumps(obj))


class _QuietHTTP(SimpleHTTPRequestHandler):
    def log_message(self, *args):
        pass

    def log_request(self, code='-', size='-'):
        path = self.path.split('?')[0].split('#')[0]
        if path.endswith('/') or path.endswith('.html'):
            ip = (self.headers.get('X-Forwarded-For') or
                  self.client_address[0]).split(',')[0].strip()
            log_stat('page', p=path, u=anon(ip))


def _argval(flag, default):
    args = sys.argv[1:]
    return args[args.index(flag) + 1] if flag in args else default


def main():
    global TOKENS
    os.chdir(os.path.dirname(os.path.abspath(__file__)))
    bind = _argval('--bind', '0.0.0.0')          # e.g. a tailnet IP behind a reverse proxy
    http_port = int(_argval('--http-port', HTTP_PORT))
    tokens_file = _argval('--tokens', None)      # one access code per line, # comments ok
    if tokens_file:
        with open(tokens_file, encoding='utf-8') as f:
            TOKENS = {l.strip() for l in f if l.strip() and not l.startswith('#')}
        print('relay hardened: %d access codes loaded' % len(TOKENS))
    global STATS
    STATS = _argval('--stats', None)             # JSONL metrics: pages, joins, session lengths
    if STATS:
        print('metrics -> %s (anonymous: daily-salted hashes, no IPs stored)' % STATS)

    socketserver.ThreadingTCPServer.allow_reuse_address = True
    ThreadingHTTPServer.allow_reuse_address = True
    httpd = ThreadingHTTPServer((bind, http_port), _QuietHTTP)
    wsd = socketserver.ThreadingTCPServer((bind, WS_PORT), WSHandler)
    httpd.daemon_threads = True
    wsd.daemon_threads = True
    _SERVERS.extend([httpd, wsd])

    ip = bind if bind not in ('0.0.0.0', '') else local_ip()
    print('backnforth running:')
    print('  Therapist page:  http://%s:%d/' % (ip, http_port))
    print('  Session relay:   ws://%s:%d' % (ip, WS_PORT))
    print('Press Ctrl+C to stop.')

    threading.Thread(target=wsd.serve_forever, daemon=True).start()
    if '--no-browser' not in sys.argv:
        try:
            webbrowser.open('http://%s:%d/easy.html' % (ip, http_port))
        except Exception:
            pass
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        print('\nStopped.')


if __name__ == '__main__':
    main()