"""
Lynk – Real‑time event engine with native HTTP routing
Pure Python, standard library only.

Features (now with AUTO protocol – TCP+UDP concurrently):
- HTTP routing with path parameters and method shortcuts
- Route groups with prefixes and middleware
- WebSocket event handling (RFC6455) with fragmentation and binary support
- UDP datagram server (protocol="UDP") – each datagram is a JSON message with a "path"
- **AUTO mode**: runs TCP (HTTP/WebSocket) and UDP on the same port simultaneously
- Client‑ID token for UDP rate limiting (falls back to IP:port)
- Max payload size enforcement for UDP
- Middleware and room-based pub/sub
- Static file serving with streaming
- Template rendering
- Background tasks and scheduler (cron‑like)
- Connection limits, origin validation, rate limiting
- Automatic heartbeat (ping/pong)
- Graceful shutdown with signal handling
- CORS support
- Plugin system
- CLI entry point
- Clean response API (redirect, json, file streaming, abort)
- Integrated logging to soketDB (HTTP, WebSocket, runtime) with auto‑sync or manual control
- Distributed database query API (run queries on any registered soketDB instance)
- Enhanced file serving: custom headers, attachment downloads, cache control, last-modified
"""

import asyncio
import hashlib
import json
import logging
import mimetypes
import os
import re
import signal
import socket
import struct
import time
import uuid
from datetime import datetime, timezone
import urllib.parse
from collections import deque, defaultdict
from contextvars import ContextVar
from typing import Any, Awaitable, Callable, Dict, List, Optional, Set, Tuple, Union, Pattern
# if soketdb is inside lynkio/
from lynkio.soketdb import database, env   # <-- soketDB integration

# JAVASCRIPT client
CLIENT_JS = """
//Lynkio javascript CLIENT
class LynkClient {
    constructor(url) {
        this.url = url;
        this.ws = null;
        this.handlers = {};
        this.binaryHandlers = [];
        this.connectionPromise = null;
        this._connectResolve = null;
        this._connectReject = null;
        this.reconnectAttempts = 0;
        this.maxReconnectAttempts = 10;
        this.reconnectDelay = 1000; // initial delay 1 second
        this.shouldReconnect = true;
    }

    connect() {
        // If already connecting or connected, return existing promise
        if (this.connectionPromise) return this.connectionPromise;

        this.connectionPromise = new Promise((resolve, reject) => {
            this._connectResolve = resolve;
            this._connectReject = reject;
        });

        this.ws = new WebSocket(this.url);
        this.ws.binaryType = 'arraybuffer';

        this.ws.onmessage = (event) => {
            if (event.data instanceof ArrayBuffer) {
                const data = new Uint8Array(event.data);
                this.binaryHandlers.forEach(handler => handler(data));
            } else {
                try {
                    const msg = JSON.parse(event.data);
                    const { event: evt, data } = msg;
                    if (this.handlers[evt]) {
                        this.handlers[evt](data);
                    }
                } catch (e) {
                    console.error('Failed to parse message', e);
                }
            }
        };

        this.ws.onopen = () => {
            console.log('WebSocket connected');
            this.reconnectAttempts = 0;
            if (this._connectResolve) {
                this._connectResolve();
                this._connectResolve = null;
                this._connectReject = null;
            }
        };

        this.ws.onclose = (event) => {
            console.log(`WebSocket closed (code: ${event.code})`);
            // If connection was never established, reject the promise
            if (this._connectReject) {
                this._connectReject(new Error('Connection closed before open'));
                this._connectResolve = null;
                this._connectReject = null;
            }
            this.connectionPromise = null;
            this.ws = null;

            if (this.shouldReconnect && this.reconnectAttempts < this.maxReconnectAttempts) {
                const delay = this.reconnectDelay * Math.pow(2, this.reconnectAttempts);
                this.reconnectAttempts++;
                console.log(`Reconnecting in ${delay}ms (attempt ${this.reconnectAttempts})`);
                setTimeout(() => this.connect(), delay);
            } else if (this.reconnectAttempts >= this.maxReconnectAttempts) {
                console.log('Max reconnect attempts reached');
            }
        };

        this.ws.onerror = (err) => {
            console.error('WebSocket error', err);
            // If connection promise is still pending, reject it
            if (this._connectReject) {
                this._connectReject(err);
                this._connectResolve = null;
                this._connectReject = null;
            }
        };

        return this.connectionPromise;
    }

    on(event, callback) {
        this.handlers[event] = callback;
    }

    onBinary(callback) {
        this.binaryHandlers.push(callback);
    }

    emit(event, data) {
        if (this.ws && this.ws.readyState === WebSocket.OPEN) {
            this.ws.send(JSON.stringify({ event, data }));
        } else {
            console.warn('WebSocket not open, cannot emit', event);
        }
    }

    sendBinary(data) {
        if (this.ws && this.ws.readyState === WebSocket.OPEN) {
            this.ws.send(data);
        } else {
            console.warn('WebSocket not open, cannot send binary');
        }
    }

    /**
     * Send a message that will be routed as a UDP datagram.
     * The server must have a WebSocket handler for the "__udp" event
     * that forwards the message to its UDP router.
     * @param {string} path - The UDP route path (e.g., "/ping")
     * @param {object} data - The payload to send
     */
    sendUdp(path, data) {
        this.emit('__udp', { path, data });
    }

    joinRoom(room) {
        this.emit('join', { room });
    }

    leaveRoom(room) {
        this.emit('leave', { room });
    }

    setSession(key, value) {
        this.emit('set_session', { key, value });
    }

    close() {
        this.shouldReconnect = false; // prevent reconnection
        if (this.ws) {
            this.ws.close();
            this.ws = null;
        }
        this.connectionPromise = null;
        this._connectResolve = null;
        this._connectReject = null;
    }
}"""

# ----------------------------------------------------------------------
# Global registry of database instances
# ----------------------------------------------------------------------
_databases: Dict[str, database] = {}

# ----------------------------------------------------------------------
# Exceptions
# ----------------------------------------------------------------------
class StopProcessing(Exception):
    """Raised in middleware to stop further processing of an event."""
    pass

class WebSocketError(Exception):
    """WebSocket protocol error."""
    pass

class HTTPError(Exception):
    """HTTP error (e.g., 404, 405)."""
    def __init__(self, status_code: int, message: str = ""):
        self.status_code = status_code
        self.message = message

# ----------------------------------------------------------------------
# WebSocket frame helpers (RFC6455)
# ----------------------------------------------------------------------
def decode_frame(data: bytes) -> Tuple[bool, int, bool, int, bytes, bytes, int]:
    """
    Decode a single WebSocket frame.
    Returns:
        fin, opcode, masked, payload_length, payload, masking_key, total_consumed
    """
    if len(data) < 2:
        raise WebSocketError("Incomplete frame header")
    b1, b2 = data[0], data[1]
    fin = (b1 & 0x80) != 0
    opcode = b1 & 0x0F
    masked = (b2 & 0x80) != 0
    payload_len = b2 & 0x7F

    index = 2
    if payload_len == 126:
        if len(data) < 4:
            raise WebSocketError("Incomplete extended payload length")
        payload_len = struct.unpack("!H", data[2:4])[0]
        index = 4
    elif payload_len == 127:
        if len(data) < 10:
            raise WebSocketError("Incomplete extended payload length")
        payload_len = struct.unpack("!Q", data[2:10])[0]
        index = 10

    if masked:
        if len(data) < index + 4:
            raise WebSocketError("Missing masking key")
        masking_key = data[index:index + 4]
        index += 4
    else:
        masking_key = b""

    if len(data) < index + payload_len:
        raise WebSocketError("Incomplete payload")

    payload = data[index:index + payload_len]
    if masked:
        payload = bytes(b ^ masking_key[i % 4] for i, b in enumerate(payload))

    total_consumed = index + payload_len
    return fin, opcode, masked, payload_len, payload, masking_key, total_consumed


def encode_frame(payload: bytes, opcode: int = 0x1, fin: bool = True, mask: bool = False) -> bytes:
    """Encode a WebSocket frame."""
    b1 = (0x80 if fin else 0) | (opcode & 0x0F)
    payload_len = len(payload)
    if payload_len < 126:
        b2 = (0x80 if mask else 0) | payload_len
        header = struct.pack("!BB", b1, b2)
    elif payload_len < 65536:
        b2 = (0x80 if mask else 0) | 126
        header = struct.pack("!BBH", b1, b2, payload_len)
    else:
        b2 = (0x80 if mask else 0) | 127
        header = struct.pack("!BBQ", b1, b2, payload_len)

    if mask:
        masking_key = struct.pack("!I", uuid.uuid4().fields[0] & 0xFFFFFFFF)
        masked_payload = bytes(b ^ masking_key[i % 4] for i, b in enumerate(payload))
        return header + masking_key + masked_payload
    else:
        return header + payload


def make_handshake_accept(key: str) -> str:
    """Generate Sec‑WebSocket‑Accept header value."""
    import base64
    GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
    sha1 = hashlib.sha1((key + GUID).encode()).digest()
    return base64.b64encode(sha1).decode()


# ----------------------------------------------------------------------
# HTTP request and response
# ----------------------------------------------------------------------
class Request:
    """Parsed HTTP request with helper methods."""
    def __init__(self, method: str, path: str, headers: Dict[str, str], body: bytes = b"", client_ip: str = ""):
        self.method = method.upper()
        self.path = path
        self.headers = headers
        self.body = body
        self.client_ip = client_ip
        self._json = None
        self._form = None
        # Added for logging
        self.start_time: Optional[float] = None
        self.request_id: Optional[str] = None

    @property
    def query_params(self) -> Dict[str, str]:
        """Parse query string into dict."""
        if "?" not in self.path:
            return {}
        query = self.path.split("?", 1)[1]
        return dict(urllib.parse.parse_qsl(query))

    @property
    def cookies(self) -> Dict[str, str]:
        """Parse Cookie header."""
        cookie_str = self.headers.get("cookie", "")
        cookies = {}
        for item in cookie_str.split(";"):
            if "=" in item:
                key, val = item.strip().split("=", 1)
                cookies[key] = val
        return cookies

    async def json(self) -> Any:
        """Parse body as JSON (cached)."""
        if self._json is None:
            self._json = json.loads(self.body)
        return self._json

    async def form(self) -> Dict[str, str]:
        """Parse body as URL-encoded form."""
        if self._form is None:
            self._form = dict(urllib.parse.parse_qsl(self.body.decode()))
        return self._form

    @classmethod
    def create_minimal(cls, method: str, path: str, headers: Dict[str, str], client_ip: str) -> 'Request':
        """Create a minimal request object (e.g., for early error logging)."""
        req = cls(method, path, headers, b"", client_ip)
        req.start_time = time.time()
        req.request_id = str(uuid.uuid4())
        return req


def http_response(status_code: int, content_type: str, body: Union[str, bytes]) -> bytes:
    """Build an HTTP response."""
    if isinstance(body, str):
        body = body.encode()
    status_text = {
        101: "Switching Protocols",
        200: "OK",
        201: "Created",
        202: "Accepted",
        204: "No Content",
        302: "Found",
        400: "Bad Request",
        403: "Forbidden",
        404: "Not Found",
        405: "Method Not Allowed",
        413: "Payload Too Large",
        429: "Too Many Requests",
        500: "Internal Server Error",
        503: "Service Unavailable",
    }.get(status_code, "Unknown")
    headers = [
        f"HTTP/1.1 {status_code} {status_text}",
        f"Content-Type: {content_type}",
        f"Content-Length: {len(body)}",
        "Connection: close",  # simplified; keep-alive can be added
        "\r\n"
    ]
    return "\r\n".join(headers).encode() + body


def abort(status_code: int, message: str = ""):
    """Raise an HTTP error to be caught by the server."""
    raise HTTPError(status_code, message)


# ----------------------------------------------------------------------
# Connection class (WebSocket client)
# ----------------------------------------------------------------------
class Connection:
    """Represents a single WebSocket client."""
    def __init__(
        self,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
        client_id: Optional[str] = None
    ):
        self.reader = reader
        self.writer = writer
        self.id = client_id or str(uuid.uuid4())
        self.session: Dict[str, Any] = {}
        self.rooms: Set[str] = set()
        self._close_requested = False
        # Fragmentation support
        self._fragmented_buffer: Optional[bytearray] = None
        self._fragmented_opcode: Optional[int] = None
        # Heartbeat
        self.last_pong = time.time()
        self.last_ping = 0.0

    async def send(self, payload: Union[str, bytes], text: bool = True) -> None:
        """Send a message to this client."""
        if isinstance(payload, str):
            payload = payload.encode()
        opcode = 0x1 if text else 0x2
        frame = encode_frame(payload, opcode=opcode, mask=False)
        try:
            self.writer.write(frame)
            await self.writer.drain()
        except (ConnectionError, BrokenPipeError):
            pass

    async def ping(self) -> None:
        """Send a ping frame."""
        self.last_ping = time.time()
        ping_payload = b"ping"
        frame = encode_frame(ping_payload, opcode=0x9, mask=False)
        try:
            self.writer.write(frame)
            await self.writer.drain()
        except (ConnectionError, BrokenPipeError):
            pass

    async def close(self, code: int = 1000, reason: str = "") -> None:
        """Send a close frame and close the connection."""
        if self._close_requested:
            return
        self._close_requested = True
        payload = struct.pack("!H", code) + reason.encode()
        frame = encode_frame(payload, opcode=0x8, mask=False)
        try:
            self.writer.write(frame)
            await self.writer.drain()
        except Exception:
            pass
        self.writer.close()
        await self.writer.wait_closed()

    @property
    def closed(self) -> bool:
        return self._close_requested or self.writer.is_closing()


# ----------------------------------------------------------------------
# Response classes / helpers
# ----------------------------------------------------------------------
def redirect(location: str, status: int = 302) -> Dict[str, Any]:
    """Return a redirect response."""
    return {"_redirect": location, "_status": status}

def json_response(data: Any, status: int = 200) -> Dict[str, Any]:
    """Return a JSON response."""
    return {"_json": data, "_status": status}

class FileResponse:
    """Streaming file response with advanced headers."""
    def __init__(
        self,
        filepath: str,
        chunk_size: int = 8192,
        content_type: Optional[str] = None,
        headers: Optional[Dict[str, str]] = None,
        as_attachment: bool = False,
        attachment_filename: Optional[str] = None,
        cache_control: Optional[str] = None,
        last_modified: Optional[float] = None,
    ):
        self.filepath = filepath
        self.chunk_size = chunk_size
        self.content_type = content_type or mimetypes.guess_type(filepath)[0] or "application/octet-stream"
        self.size = os.path.getsize(filepath)
        self.headers = headers or {}
        self.as_attachment = as_attachment
        self.attachment_filename = attachment_filename or os.path.basename(filepath)
        self.cache_control = cache_control
        # Last-Modified: use provided timestamp or file's mtime
        self.last_modified = last_modified or os.path.getmtime(filepath)
        # Add Last-Modified header automatically unless overridden
        if "Last-Modified" not in self.headers:
            # RFC 7231 requires HTTP-date format (GMT)
            self.headers["Last-Modified"] = datetime.fromtimestamp(self.last_modified, tz=timezone.utc).strftime("%a, %d %b %Y %H:%M:%S GMT")
        if self.cache_control:
            self.headers["Cache-Control"] = self.cache_control
        if self.as_attachment:
            self.headers["Content-Disposition"] = f'attachment; filename="{self.attachment_filename}"'

    async def __aiter__(self):
        with open(self.filepath, "rb") as f:
            while chunk := f.read(self.chunk_size):
                yield chunk


class StreamingResponse:
    """Generic streaming response with chunked encoding."""
    def __init__(self, generator, content_type: str = "application/octet-stream"):
        self.generator = generator
        self.content_type = content_type

    async def __aiter__(self):
        async for chunk in self.generator:
            yield chunk


# ----------------------------------------------------------------------
# Route group
# ----------------------------------------------------------------------
class RouteGroup:
    """Group of routes with a common prefix and middleware."""
    def __init__(self, prefix: str, app: 'Lynk'):
        self.prefix = prefix.rstrip("/")
        self.app = app
        self._middleware: List[Callable] = []

    def use(self, middleware: Callable):
        """Add middleware to this group."""
        self._middleware.append(middleware)
        return self

    def route(self, path: str, methods: Optional[List[str]] = None):
        """Register a route within this group."""
        full_path = self.prefix + path
        def decorator(func):
            # Wrap the handler to run group middleware first
            async def wrapped(req, *args, **kwargs):
                # Run group middleware (if any)
                for mw in self._middleware:
                    # Middleware can modify request or return a response early
                    # For simplicity, we assume middleware that returns None or a response
                    # We'll implement a simple middleware that can return a response
                    resp = await mw(req)
                    if resp is not None:
                        return resp
                return await func(req, *args, **kwargs)
            self.app.route(full_path, methods)(wrapped)
            return wrapped
        return decorator

    def get(self, path: str):
        return self.route(path, methods=["GET"])

    def post(self, path: str):
        return self.route(path, methods=["POST"])

    def put(self, path: str):
        return self.route(path, methods=["PUT"])

    def delete(self, path: str):
        return self.route(path, methods=["DELETE"])

    def patch(self, path: str):
        return self.route(path, methods=["PATCH"])


# ----------------------------------------------------------------------
# CORS middleware
# ----------------------------------------------------------------------
def cors_middleware(allowed_origins: Optional[List[str]] = None, allow_credentials: bool = False):
    """Return a middleware that adds CORS headers and handles preflight."""
    allowed = allowed_origins or ["*"]

    async def middleware(req):
        # Handle preflight OPTIONS request
        if req.method == "OPTIONS":
            resp_headers = {
                "Access-Control-Allow-Origin": req.headers.get("origin", "*") if "*" in allowed else allowed[0],
                "Access-Control-Allow-Methods": "GET, POST, PUT, DELETE, PATCH, OPTIONS",
                "Access-Control-Allow-Headers": req.headers.get("access-control-request-headers", "*"),
                "Access-Control-Max-Age": "86400",
            }
            if allow_credentials:
                resp_headers["Access-Control-Allow-Credentials"] = "true"
            # Return a 200 response with no body
            response = "\r\n".join([f"HTTP/1.1 200 OK"] + [f"{k}: {v}" for k, v in resp_headers.items()] + ["", ""])
            return response.encode()
        # For non-OPTIONS, store origin for later use in response
        req._cors_origin = req.headers.get("origin", "*")
        req._cors_credentials = allow_credentials
        return None
    return middleware

# ======================================================================
# UDP Protocol Handler
# ======================================================================
class UDPProtocol(asyncio.DatagramProtocol):
    def __init__(self, app: 'Lynk'):
        self.app = app
        self.transport = None

    def connection_made(self, transport: asyncio.DatagramTransport):
        self.transport = transport
        self.app._logger.info("UDP server listening")

    def datagram_received(self, data: bytes, addr: Tuple[str, int]):
        client_ip, client_port = addr
        asyncio.create_task(self.app._handle_udp_datagram(data, addr, client_ip, client_port))

    def error_received(self, exc):
        self.app._logger.error(f"UDP error: {exc}")

    def connection_lost(self, exc):
        self.app._logger.info("UDP server stopped")

# ----------------------------------------------------------------------
# Lynk – main engine (HTTP + WebSocket + UDP + AUTO)
# ----------------------------------------------------------------------
class Lynk:
    """Main event engine. Manages WebSocket clients, HTTP routes, UDP datagrams, and AUTO mode."""

    def __init__(
        self,
        host: str = "0.0.0.0",
        port: int = 8765,
        protocol: str = "TCP",                    # "TCP", "UDP", or "AUTO"
        max_payload_size: int = 256 * 1024,       # 256 KiB per frame / datagram
        max_message_size: int = 1024 * 1024,      # 1 MiB total fragmented message
        max_body_size: int = 1024 * 1024,         # 1 MiB HTTP body
        room_batch_size: int = 100,
        max_connections: Optional[int] = None,
        allowed_origins: Optional[List[str]] = None,
        rate_limit: Optional[int] = None,          # messages per second per client/UDP token
        enable_keep_alive: bool = False,
        debug: bool = False,
        enable_database: bool = False,
        database_config: Optional[dict] = None,
        serve_client: bool = False,
        client_path: str = "/lynkio/client.js"
    ):
        self.protocol = protocol.upper()           # "TCP", "UDP", "AUTO"
        self.serve_client = serve_client
        self.client_path = client_path
        self.host = host
        self.enable_database = enable_database
        self.database_config = database_config
        self.port = port
        self.max_payload_size = max_payload_size
        self.max_message_size = max_message_size
        self.max_body_size = max_body_size
        self.room_batch_size = room_batch_size
        self.max_connections = max_connections
        self.allowed_origins = allowed_origins or []
        self.rate_limit = rate_limit
        self.enable_keep_alive = enable_keep_alive
        self.debug = debug

        # WebSocket (only used in TCP/AUTO mode)
        self._clients: Dict[str, Connection] = {}
        self._rooms: Dict[str, Set[str]] = {}
        self._handlers: Dict[str, Callable[[Connection, Any], Awaitable[None]]] = {}
        self._binary_handlers: List[Callable[[Connection, bytes], Awaitable[None]]] = []
        self._internal_handlers: Dict[
            str, List[Callable[[Optional[Connection], Optional[Any]], Awaitable[None]]]
        ] = {}
        self._middleware: List[
            Callable[[Connection, str, Any], Awaitable[Optional[Any]]]
        ] = []

        # HTTP/UDP routes: list of (pattern, handler, methods)
        self._http_routes: List[Tuple[Pattern, Callable, Set[str]]] = []

        # HTTP middleware (also used for UDP)
        self._http_middleware: List[Callable] = []

        # Background tasks
        self._background_tasks: List[Callable[[], Awaitable[None]]] = []
        self._scheduled_tasks: List[Tuple[float, Callable]] = []

        # Rate limiting per client (sliding window) – keys can be IP:port or client_id (UDP)
        self._client_msg_times: Dict[str, deque] = defaultdict(lambda: deque(maxlen=rate_limit or 0))

        # Server state
        self._tcp_server: Optional[asyncio.Server] = None          # renamed for clarity
        self._udp_transport: Optional[asyncio.DatagramTransport] = None
        self._udp_protocol: Optional[UDPProtocol] = None
        self._loop: Optional[asyncio.AbstractEventLoop] = None
        self._running = False
        self._shutdown_event = asyncio.Event()
        self._logger = logging.getLogger("Lynk")
        if debug:
            self._logger.setLevel(logging.DEBUG)

        # Request ID context
        self.request_id_ctx: ContextVar[str] = ContextVar('request_id', default='')

        # CORS settings (global)
        self.cors_allowed_origins = []
        self.cors_allow_credentials = False

        # Logging integration (soketDB)
        self._db = None
        self.auto_sync_log = False
        self.loop = None

        if self.serve_client:
            self._add_client_route()

    def _add_client_route(self) -> None:
        pattern = re.compile(f"^{re.escape(self.client_path)}$")
        async def client_handler(req: Request):
            return (CLIENT_JS, "application/javascript")
        self._http_routes.append((pattern, client_handler, {"GET"}))
        self._logger.info(f"Serving built‑in lynkio client.js at {self.client_path}")

    # ------------------------------------------------------------------
    # DATABASE WRAPPER
    # ------------------------------------------------------------------
    def create_database(self, name: str = "lynkio_test_db",
                        create_log_table: bool = False,
                        auto_sync_log: bool = False):
        """Create a soketDB instance and optionally set up log tables."""
        if not self.enable_database:
            raise Exception("Database is not enabled.")

        config = self.database_config or {
            'primary_storage': 'local',
            'backup_enabled': True,
            'auto_backup_hours': 24,
            'query_cache_enabled': True,
            'auto_sync': True,
            'google_drive_enabled': False,
            'huggingface_enabled': False,
            'aws_s3_enabled': False,
            'dropbox_enabled': False
        }

        try:
            db = database(name, config)
            if db:
                self._db = db
                self.auto_sync_log = auto_sync_log
                _databases[name] = db          # store globally for distributed queries

                if create_log_table:
                    # Create the three log tables (soketDB: only column names, no types)
                    db.execute("""
                        CREATE TABLE IF NOT EXISTS wss_logs (
                            id, timestamp, client_id, direction, event, data, size, opcode
                        )
                    """)
                    db.execute("""
                        CREATE TABLE IF NOT EXISTS http_logs (
                            id, timestamp, method, path, status_code, client_ip, user_agent, response_time, request_id
                        )
                    """)
                    db.execute("""
                        CREATE TABLE IF NOT EXISTS runtime_logs (
                            id, timestamp, level, message, source
                        )
                    """)
                return db
            return db
        except Exception as e:
            print(e)
            return None

    # ------------------------------------------------------------------
    # Distributed query API (now accepts optional params)
    # ------------------------------------------------------------------
    async def query_database(self, db_name: str, query: str, params: Optional[Tuple] = None) -> Any:
        """
        Execute a query on a named database (from the global _databases registry) asynchronously.
        Returns the result of db.execute(query, params).  
        If parameters are provided, they must be a tuple of values matching $n placeholders in the query.
        """
        db = _databases.get(db_name)
        if not db:
            raise ValueError(f"Database '{db_name}' not found. Available: {list(_databases.keys())}")
        loop = self._loop or asyncio.get_running_loop()
        if params is None:
            return await loop.run_in_executor(None, db.execute, query)
        else:
            return await loop.run_in_executor(None, db.execute, query, params)

    # ------------------------------------------------------------------
    # Helper to run DB queries in executor (non‑blocking, for internal use)
    # ------------------------------------------------------------------
    async def _db_execute(self, query: str, params=None):
        """Run a soketDB execute in a thread executor (for the primary logging DB)."""
        if not self._db:
            return
        loop = self._loop or asyncio.get_running_loop()
        if params:
          await loop.run_in_executor(None, self._db.execute, query, params)
        else:
          await loop.run_in_executor(None, self._db.execute, query)

    # ------------------------------------------------------------------
    # Core logging methods (auto‑sync) – adapted to soketDB syntax
    # ------------------------------------------------------------------
    async def _log_http(self, req: Request, status_code: int):
        """Log an HTTP request/response: print to console and optionally to database."""
        # Always print to console (info for 2xx/3xx, warning for errors)
        response_time = time.time() - req.start_time
        client_ip = req.client_ip or "unknown"
        log_msg = f"HTTP {req.method} {req.path} -> {status_code} ({response_time:.3f}s) from {client_ip}"
        if 200 <= status_code < 400:
            self._logger.info(log_msg)
        else:
            self._logger.warning(log_msg)

        # Database logging (if enabled)
        if not self._db or not self.auto_sync_log:
            return
        user_agent = req.headers.get('user-agent', '')
        log_id = str(uuid.uuid4())
        timestamp = datetime.now().isoformat()

        query = """
            INSERT INTO http_logs (id, timestamp, method, path, status_code, client_ip, user_agent, response_time, request_id)
            VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
        """
        params = (
            log_id,
            timestamp,
            req.method,
            req.path,
            status_code,
            req.client_ip,
            user_agent,
            response_time,
            req.request_id
        )
        await self._db_execute(query, params)

    async def _log_websocket(self, client_id: str, direction: str,
                             event: str = None, data: Any = None,
                             opcode: int = None, size: int = 0):
        """Log a WebSocket message/event if auto_sync_log is True."""
        if not self._db or not self.auto_sync_log:
            return
        data_str = None
        if data is not None:
            if isinstance(data, (dict, list)):
                data_str = json.dumps(data)
            elif isinstance(data, bytes):
                data_str = data.hex()  # or truncate for large binaries
            else:
                data_str = str(data)

        log_id = str(uuid.uuid4())
        timestamp = datetime.now().isoformat()

        query = """
            INSERT INTO wss_logs (id, timestamp, client_id, direction, event, data, size, opcode)
            VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
        """
        params = (
            log_id,
            timestamp,
            client_id,
            direction,
            event,
            data_str,
            size,
            opcode
        )
        await self._db_execute(query, params)

    async def _log_runtime(self, level: str, message: str, source: str = "server"):
        """Log a runtime event if auto_sync_log is True."""
        if not self._db or not self.auto_sync_log:
            return
        log_id = str(uuid.uuid4())
        timestamp = datetime.now().isoformat()

        query = """
            INSERT INTO runtime_logs (id, timestamp, level, message, source)
            VALUES ($1, $2, $3, $4, $5)
        """
        params = (
            log_id,
            timestamp,
            level,
            message,
            source
        )
        await self._db_execute(query, params)

    # ------------------------------------------------------------------
    # Public manual logging helper – now parameterised
    # ------------------------------------------------------------------
    async def add_log(self, table: str, **kwargs):
        """
        Manually insert a log entry into the specified table.
        Table must be one of: 'http', 'wss', 'runtime'.
        Keyword arguments must match the table's columns (except id/timestamp are generated automatically).
        """
        if not self._db:
            return
        table_map = {
            'http': 'http_logs',
            'wss': 'wss_logs',
            'runtime': 'runtime_logs'
        }
        table_name = table_map.get(table)
        if not table_name:
            raise ValueError("Table must be 'http', 'wss', or 'runtime'")

        # Auto‑generate id and timestamp
        kwargs['id'] = str(uuid.uuid4())
        kwargs['timestamp'] = datetime.now().isoformat()

        # Build parameterised INSERT dynamically
        columns = list(kwargs.keys())
        placeholders = [f'${i+1}' for i in range(len(columns))]
        query = f"INSERT INTO {table_name} ({', '.join(columns)}) VALUES ({', '.join(placeholders)})"
        params = tuple(kwargs[col] for col in columns)

        await self._db_execute(query, params)

    # ------------------------------------------------------------------
    # Public API: CORS
    # ------------------------------------------------------------------
    def enable_cors(self, allowed_origins: Optional[List[str]] = None, allow_credentials: bool = False):
        """Enable CORS globally."""
        self.cors_allowed_origins = allowed_origins or ["*"]
        self.cors_allow_credentials = allow_credentials
        # Add the CORS middleware
        mw = cors_middleware(allowed_origins, allow_credentials)
        self._http_middleware.append(mw)

    # ------------------------------------------------------------------
    # Public API: Route groups
    # ------------------------------------------------------------------
    def group(self, prefix: str) -> RouteGroup:
        """Create a route group with a common prefix."""
        return RouteGroup(prefix, self)

    # ------------------------------------------------------------------
    # Public API: WebSocket decorators
    # ------------------------------------------------------------------
    def on(self, event: str):
        """Register an event handler for JSON messages."""
        def decorator(func: Callable[[Connection, Any], Awaitable[None]]):
            self._handlers[event] = func
            return func
        return decorator

    def on_binary(self, func: Callable[[Connection, bytes], Awaitable[None]]):
        """Register a handler for binary WebSocket messages."""
        self._binary_handlers.append(func)
        return func

    def on_internal(self, event: str):
        """Register an internal event handler."""
        def decorator(func: Callable[[Optional[Connection], Optional[Any]], Awaitable[None]]):
            self._internal_handlers.setdefault(event, []).append(func)
            return func
        return decorator

    def middleware(self, func):
        """Register middleware for WebSocket."""
        self._middleware.append(func)
        return func

    # ------------------------------------------------------------------
    # Public API: HTTP decorators (with method shortcuts)
    # ------------------------------------------------------------------
    def route(self, path: str, methods: Optional[List[str]] = None):
        if methods is None:
            methods = ["GET"]
        methods = {m.upper() for m in methods}
        pattern = re.compile(f"^{path}$")

        def decorator(func: Callable[[Request], Awaitable[Any]]):
            self._http_routes.append((pattern, func, methods))
            return func
        return decorator

    def get(self, path: str):
        return self.route(path, methods=["GET"])

    def post(self, path: str):
        return self.route(path, methods=["POST"])

    def put(self, path: str):
        return self.route(path, methods=["PUT"])

    def delete(self, path: str):
        return self.route(path, methods=["DELETE"])

    def patch(self, path: str):
        return self.route(path, methods=["PATCH"])

    def udp(self, path: str):
        """Register a UDP route (method set to {'UDP'})."""
        return self.route(path, methods=["UDP"])

    # ------------------------------------------------------------------
    # Public API: Unified route (HTTP + WebSocket)
    # ------------------------------------------------------------------
    def both(self, path: str):
        """Register a handler for both HTTP GET and WebSocket events on the same path."""
        # For WebSocket, we'll use the path as the event name
        def decorator(func):
            # HTTP part
            @self.get(path)
            async def http_handler(req):
                return await func(req)
            # WebSocket part
            @self.on(path)
            async def ws_handler(client, data):
                # We need to adapt: the function may expect (req) or (client, data)
                # This is tricky; we'll assume the function takes (client, data) for WebSocket
                await func(client, data)
            return func
        return decorator

    # ------------------------------------------------------------------
    # Static file serving
    # ------------------------------------------------------------------
    def static(self, prefix: str, directory: str):
        """Serve static files from a directory under a URL prefix."""
        prefix = prefix.rstrip("/")
        pattern = re.compile(f"^{prefix}/(?P<filepath>.*)$")
        base_abs = os.path.abspath(directory)  # absolute path to static folder

        async def serve_static(req: Request, filepath: str):
            # Security: prevent directory traversal
            full_path = os.path.abspath(os.path.join(directory, filepath))
            if not full_path.startswith(base_abs):
                abort(403, "Forbidden")
            if not os.path.exists(full_path) or os.path.isdir(full_path):
                abort(404, "Not Found")
            # Use enhanced send_file with default settings
            return send_file(full_path, base_dir="")

        self._http_routes.append((pattern, serve_static, {"GET"}))

    # ------------------------------------------------------------------
    # Background tasks and scheduler
    # ------------------------------------------------------------------
    def task(self, func: Callable[[], Awaitable[None]]):
        """Register a background coroutine to run when the server starts."""
        self._background_tasks.append(func)
        return func

    def schedule(self, interval: float):
        """Decorator to run a coroutine periodically (in seconds)."""
        def decorator(func: Callable[[], Awaitable[None]]):
            self._scheduled_tasks.append((interval, func))
            return func
        return decorator

    # ------------------------------------------------------------------
    # Plugin system
    # ------------------------------------------------------------------
    def use(self, plugin):
        """Register a plugin (a callable that takes the app instance)."""
        plugin(self)

    # ------------------------------------------------------------------
    # Public API: WebSocket actions
    # ------------------------------------------------------------------
    async def emit(self, event: str, data: Any, client_id: Optional[str] = None) -> None:
        """Send to one client or broadcast to all."""
        if client_id is not None:
            client = self._clients.get(client_id)
            if client and not client.closed:
                self._logger.debug(f"Emitting {event} to client {client_id}")
                message = json.dumps({"event": event, "data": data})
                await client.send(message)
                # Log outgoing message
                await self._log_websocket(client_id, 'out', event=event,
                                          data=data, size=len(message))
            else:
                self._logger.debug(f"Client {client_id} not found or closed, cannot emit {event}")
        else:
            self._logger.debug(f"Broadcasting {event} to all clients")
            await self.emit_to_all_except(event, data, exclude_ids=[])

    async def emit_to_all_except(self, event: str, data: Any, exclude_ids: List[str]) -> None:
        """Broadcast to all clients except those listed in exclude_ids."""
        message = json.dumps({"event": event, "data": data})
        tasks = []
        for cid, c in self._clients.items():
            if not c.closed and cid not in exclude_ids:
                tasks.append(c.send(message))
                # Log each outgoing message (could be many; consider rate‑limiting)
                await self._log_websocket(cid, 'out', event=event,
                                          data=data, size=len(message))
        if tasks:
            self._logger.debug(f"Broadcasting {event} to {len(tasks)} clients (excluding {exclude_ids})")
            await asyncio.gather(*tasks, return_exceptions=True)

    async def emit_to_room(
        self, room: str, event: str, data: Any, exclude: Optional[str] = None
    ) -> None:
        """Send an event to all clients in a room (chunked to avoid loop flooding)."""
        if room not in self._rooms:
            self._logger.debug(f"Room {room} does not exist, cannot emit {event}")
            return
        message = json.dumps({"event": event, "data": data})
        client_ids = [cid for cid in self._rooms[room] if cid != exclude]
        if not client_ids:
            self._logger.debug(f"No clients in room {room} to emit {event} (exclude={exclude})")
            return

        self._logger.debug(f"Emitting {event} to room {room} for {len(client_ids)} clients")
        for i in range(0, len(client_ids), self.room_batch_size):
            batch = client_ids[i:i + self.room_batch_size]
            tasks = []
            for cid in batch:
                client = self._clients.get(cid)
                if client and not client.closed:
                    tasks.append(client.send(message))
                    # Log outgoing message
                    await self._log_websocket(cid, 'out', event=event,
                                              data=data, size=len(message))
            if tasks:
                await asyncio.gather(*tasks, return_exceptions=True)

    def join_room(self, client_id: str, room: str) -> None:
        """Add a client to a room."""
        client = self._clients.get(client_id)
        if not client or client.closed:
            self._logger.debug(f"Cannot join room {room}: client {client_id} not found or closed")
            return
        self._rooms.setdefault(room, set()).add(client_id)
        client.rooms.add(room)
        self._logger.debug(f"Client {client_id} joined room {room}")
        self._dispatch_internal("subscribe", client, {"room": room})

    def leave_room(self, client_id: str, room: str) -> None:
        """Remove a client from a room."""
        client = self._clients.get(client_id)
        if not client:
            return
        if room in self._rooms:
            self._rooms[room].discard(client_id)
            if not self._rooms[room]:
                del self._rooms[room]
                self._logger.debug(f"Room {room} deleted (empty)")
        client.rooms.discard(room)
        self._logger.debug(f"Client {client_id} left room {room}")
        self._dispatch_internal("unsubscribe", client, {"room": room})

    def get_room_clients(self, room: str) -> Set[str]:
        """Return a copy of the set of client IDs in a room."""
        return self._rooms.get(room, set()).copy()

    # ------------------------------------------------------------------
    # Internal event dispatching
    # ------------------------------------------------------------------
    def _dispatch_internal(
        self, event: str, client: Optional[Connection] = None, data: Optional[Any] = None
    ) -> None:
        """Schedule internal event handlers (fire‑and‑forget)."""
        if event not in self._internal_handlers:
            return
        self._logger.debug(f"Dispatching internal event {event} for client {client.id if client else 'system'}")
        for handler in self._internal_handlers[event]:
            asyncio.create_task(handler(client, data))

    # ------------------------------------------------------------------
    # Server lifecycle (AUTO mode starts both TCP and UDP)
    # ------------------------------------------------------------------
    async def start(self) -> None:
        """Start the server based on the selected protocol."""
        self._loop = asyncio.get_running_loop()
        self._running = True

        # Start TCP server if protocol is TCP or AUTO
        if self.protocol in ("TCP", "AUTO"):
            self._tcp_server = await asyncio.start_server(
                self._handle_connection, self.host, self.port
            )
            self._logger.info(f"Lynk TCP server listening on {self.host}:{self.port} (HTTP/WebSocket)")

        # Start UDP server if protocol is UDP or AUTO
        if self.protocol in ("UDP", "AUTO"):
            self._udp_protocol = UDPProtocol(self)
            self._udp_transport, _ = await self._loop.create_datagram_endpoint(
                lambda: self._udp_protocol,
                local_addr=(self.host, self.port)
            )
            self._logger.info(f"Lynk UDP server listening on {self.host}:{self.port}")

        # Log runtime start
        await self._log_runtime('INFO', f'Server started on {self.host}:{self.port} ({self.protocol})', 'server')

        # Start background tasks
        for task in self._background_tasks:
            asyncio.create_task(task())

        # Start scheduled tasks
        for interval, func in self._scheduled_tasks:
            asyncio.create_task(self._run_scheduled(interval, func))

        # Start heartbeat (only if TCP is active)
        if self.protocol in ("TCP", "AUTO"):
            asyncio.create_task(self._heartbeat())

    async def _run_scheduled(self, interval: float, func: Callable):
        """Run a scheduled task periodically."""
        while self._running:
            await asyncio.sleep(interval)
            try:
                await func()
            except Exception:
                self._logger.exception("Scheduled task error")
                await self._log_runtime('ERROR', f'Scheduled task {func.__name__} failed', 'scheduler')

    async def stop(self) -> None:
        """Stop all servers gracefully."""
        self._running = False
        self._shutdown_event.set()

        # Close TCP server if it exists
        if self._tcp_server:
            self._tcp_server.close()
            await self._tcp_server.wait_closed()
            # Close all WebSocket clients
            close_tasks = [client.close(1001, "Server shutting down") for client in self._clients.values()]
            if close_tasks:
                await asyncio.gather(*close_tasks, return_exceptions=True)
            self._clients.clear()
            self._rooms.clear()

        # Close UDP transport if it exists
        if self._udp_transport:
            self._udp_transport.close()

        self._logger.info("Lynk stopped")
        await self._log_runtime('INFO', 'Server stopped', 'server')

    def run(self) -> None:
        """Run the server until interrupted (blocking call)."""
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)

        # Signal handlers for graceful shutdown
        for sig in (signal.SIGINT, signal.SIGTERM):
            loop.add_signal_handler(sig, lambda: asyncio.create_task(self.stop()))

        try:
            loop.run_until_complete(self._run_forever())
        finally:
            loop.close()

    async def _run_forever(self) -> None:
        await self.start()
        try:
            await self._shutdown_event.wait()
        except asyncio.CancelledError:
            pass
        finally:
            await self.stop()

    # ------------------------------------------------------------------
    # Heartbeat
    # ------------------------------------------------------------------
    async def _heartbeat(self):
        """Periodically ping clients and close unresponsive ones."""
        while self._running:
            await asyncio.sleep(30)  # ping interval
            now = time.time()
            for client in list(self._clients.values()):
                if client.closed:
                    continue
                # If no pong for 60 seconds, close
                if now - client.last_pong > 60:
                    self._logger.debug(f"Client {client.id} timed out")
                    await client.close(1001, "Going away")
                else:
                    await client.ping()

    # ------------------------------------------------------------------
    # Connection handling (HTTP / WebSocket dispatcher)
    # ------------------------------------------------------------------
    async def _handle_connection(
        self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter
    ) -> None:
        """Handle a new client connection – either HTTP or WebSocket upgrade."""
        # Check connection limit (create minimal request for logging)
        if self.max_connections and len(self._clients) >= self.max_connections:
            req = Request.create_minimal("UNKNOWN", "/", {}, "")
            await self._send_http_error(writer, 503, "Server busy", req)
            return

        # Get client IP
        client_ip = writer.get_extra_info('peername')
        if client_ip:
            client_ip = client_ip[0]
        else:
            client_ip = ""

        try:
            # Read request line
            request_line = await reader.readline()
            if not request_line:
                writer.close()
                await writer.wait_closed()
                return

            # Parse request line
            parts = request_line.decode().strip().split()
            if len(parts) < 3:
                req = Request.create_minimal("UNKNOWN", "/", {}, client_ip)
                await self._send_http_error(writer, 400, "Bad Request", req)
                return
            method, path, version = parts[0], parts[1], parts[2]

            # Read headers
            headers = {}
            while True:
                line = await reader.readline()
                if line == b"\r\n":
                    break
                if not line:
                    req = Request.create_minimal(method, path, headers, client_ip)
                    await self._send_http_error(writer, 400, "Bad Request", req)
                    return
                try:
                    key, value = line.decode().strip().split(":", 1)
                    headers[key.strip().lower()] = value.strip()
                except ValueError:
                    req = Request.create_minimal(method, path, headers, client_ip)
                    await self._send_http_error(writer, 400, "Malformed header", req)
                    return

            # Create minimal request now (body not read yet)
            req = Request.create_minimal(method, path, headers, client_ip)

            # Check for WebSocket upgrade
            if headers.get("upgrade", "").lower() == "websocket":
                # Origin validation
                origin = headers.get("origin", "")
                if self.allowed_origins and origin not in self.allowed_origins:
                    await self._send_http_error(writer, 403, "Origin not allowed", req)
                    return
                # Handle WebSocket handshake and then messages
                try:
                    await self._websocket_handshake(reader, writer, headers, req)
                except Exception as e:
                    self._logger.warning(f"WebSocket handshake failed: {e}")
                    writer.close()
                    await writer.wait_closed()
                return

            # Handle HTTP request
            content_length = int(headers.get("content-length", 0))
            if content_length > self.max_body_size:
                await self._send_http_error(writer, 413, "Request entity too large", req)
                return
            body = await reader.read(content_length) if content_length > 0 else b""
            req.body = body  # update the request body

            self.request_id_ctx.set(req.request_id)

            # Run HTTP middleware
            for mw in self._http_middleware:
                try:
                    resp = await mw(req)
                    if resp is not None:
                        # If middleware returns a response, send it and stop
                        if isinstance(resp, bytes):
                            writer.write(resp)
                            await writer.drain()
                        else:
                            await self._send_http_response(writer, resp, req)
                        return
                except Exception as e:
                    self._logger.exception("HTTP middleware error")
                    await self._send_http_error(writer, 500, "Middleware error", req)
                    return

            # Find matching route
            route_path = path.split("?", 1)[0]
            for pattern, handler, methods in self._http_routes:
                match = pattern.match(route_path)
                if match:
                    if method not in methods:
                        await self._send_http_error(writer, 405, f"Method {method} not allowed", req)
                        return
                    kwargs = match.groupdict()
                    try:
                        result = await handler(req, **kwargs)
                        await self._send_http_response(writer, result, req)
                    except HTTPError as e:
                        await self._send_http_error(writer, e.status_code, e.message, req)
                    except Exception as e:
                        self._logger.exception(f"HTTP handler error for {method} {path}")
                        await self._send_http_error(writer, 500, "Internal Server Error", req)
                    break
            else:
                await self._send_http_error(writer, 404, f"Not Found: {route_path}", req)

        except Exception as e:
            self._logger.exception("Unhandled error in connection handler")
            if 'req' not in locals():
                req = Request.create_minimal("UNKNOWN", "/", {}, client_ip)
            await self._send_http_error(writer, 500, "Internal Server Error", req)
        finally:
            # If keep-alive is not enabled, close the connection
            if not self.enable_keep_alive or headers.get("connection", "").lower() != "keep-alive":
                try:
                    writer.close()
                    await writer.wait_closed()
                except:
                    pass

    async def _websocket_handshake(
        self,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
        headers: Dict[str, str],
        req: Request
    ) -> None:
        """Perform WebSocket handshake and then enter message loop."""
        key = headers.get("sec-websocket-key")
        if not key:
            raise WebSocketError("Missing Sec-WebSocket-Key")
        accept = make_handshake_accept(key)
        response = (
            "HTTP/1.1 101 Switching Protocols\r\n"
            "Upgrade: websocket\r\n"
            "Connection: Upgrade\r\n"
            f"Sec-WebSocket-Accept: {accept}\r\n"
            "\r\n"
        )
        writer.write(response.encode())
        await writer.drain()

        # Log the successful handshake (status 101)
        await self._log_http(req, 101)

        self._logger.debug("WebSocket handshake successful")

        # Create client and start message loop
        client = Connection(reader, writer)
        self._clients[client.id] = client
        self._logger.info(f"WebSocket client connected: {client.id}")
        await self._log_websocket(client.id, 'in', event='connect')
        self._dispatch_internal("connect", client)

        try:
            await self._read_websocket_messages(client)
        except (ConnectionError, asyncio.CancelledError):
            pass
        finally:
            for room in list(client.rooms):
                self.leave_room(client.id, room)
            if client.id in self._clients:
                del self._clients[client.id]
            self._logger.info(f"WebSocket client disconnected: {client.id}")
            await self._log_websocket(client.id, 'out', event='disconnect')
            self._dispatch_internal("disconnect", client)
            if not client.closed:
                await client.close()

    async def _read_websocket_messages(self, client: Connection) -> None:
        """Read and process WebSocket frames, handling fragmentation."""
        buffer = b""
        while not client.closed:
            try:
                chunk = await client.reader.read(4096)
                if not chunk:
                    break
                buffer += chunk
                while buffer:
                    try:
                        fin, opcode, masked, length, payload, key, consumed = decode_frame(buffer)
                    except WebSocketError as e:
                        self._logger.debug(f"Protocol error, need more data: {e}")
                        break

                    if length > self.max_payload_size:
                        self._logger.debug(f"Payload too large ({length} bytes), closing")
                        await client.close(1009, "Payload too large")
                        return

                    buffer = buffer[consumed:]

                    # Control frames
                    if opcode == 0x8:  # close
                        await client.close()
                        return
                    elif opcode == 0x9:  # ping
                        # Echo payload in pong (RFC6455)
                        await client.send(payload, text=False)  # pong with same payload
                        client.last_pong = time.time()
                        continue
                    elif opcode == 0xA:  # pong
                        client.last_pong = time.time()
                        continue
                    elif opcode & 0x8:  # other control frames
                        continue

                    # Data frames
                    if opcode == 0x0:  # continuation
                        if client._fragmented_buffer is None:
                            self._logger.debug("Unexpected continuation frame")
                            await client.close(1002, "Protocol error")
                            return
                        client._fragmented_buffer.extend(payload)
                        if fin:
                            # Message complete
                            complete_payload = bytes(client._fragmented_buffer)
                            op = client._fragmented_opcode
                            client._fragmented_buffer = None
                            client._fragmented_opcode = None
                            if len(complete_payload) > self.max_message_size:
                                await client.close(1009, "Message too large")
                                return
                            await self._handle_websocket_frame(client, op, complete_payload)
                    else:
                        # New message (text or binary)
                        if client._fragmented_buffer is not None:
                            self._logger.debug("New frame before previous fragmented message finished")
                            await client.close(1002, "Protocol error")
                            return
                        if not fin:
                            # Start fragmented message
                            client._fragmented_buffer = bytearray(payload)
                            client._fragmented_opcode = opcode
                        else:
                            # Single frame message
                            if len(payload) > self.max_message_size:
                                await client.close(1009, "Message too large")
                                return
                            await self._handle_websocket_frame(client, opcode, payload)
            except (ConnectionError, asyncio.CancelledError):
                break
            except Exception:
                self._logger.exception("Error in WebSocket read loop")
                break

    async def _handle_websocket_frame(self, client: Connection, opcode: int, payload: bytes) -> None:
        """Handle a complete WebSocket message (after reassembly if needed)."""
        if opcode == 0x1:  # text
            try:
                message = payload.decode()
            except UnicodeDecodeError:
                await client.close(1007, "Invalid UTF-8")
                return
            await self._handle_websocket_message(client, message)
        elif opcode == 0x2:  # binary
            await self._handle_websocket_binary(client, payload)

    async def _handle_websocket_message(self, client: Connection, message: str) -> None:
        """Handle a WebSocket text message (JSON)."""
        # Rate limiting
        if self.rate_limit:
            now = time.time()
            times = self._client_msg_times[client.id]
            # Remove messages older than 1 second
            while times and times[0] < now - 1:
                times.popleft()
            if len(times) >= self.rate_limit:
                await client.close(1008, "Rate limit exceeded")
                return
            times.append(now)

        client.session["_msg_count"] = client.session.get("_msg_count", 0) + 1
        client.session["_last_active"] = time.time()

        try:
            data = json.loads(message)
            event = data.get("event")
            payload = data.get("data", {})
        except json.JSONDecodeError:
            await client.close(1007, "Invalid JSON")
            return

        if not isinstance(event, str):
            return

        self._logger.debug(f"Received event '{event}' from {client.id} with payload: {payload}")

        # Log incoming message
        await self._log_websocket(client.id, 'in', event=event,
                                  data=payload, size=len(message))

        # Run middleware chain
        for i, middleware in enumerate(self._middleware):
            try:
                result = await middleware(client, event, payload)
                if result is not None:
                    payload = result
            except StopProcessing:
                return
            except Exception:
                self._logger.exception(f"Middleware {i} error")
                return

        handler = self._handlers.get(event)
        if handler:
            try:
                await handler(client, payload)
            except Exception:
                self._logger.exception(f"Handler error for event {event}")

        self._dispatch_internal("message", client, {"event": event, "data": payload})

    async def _handle_websocket_binary(self, client: Connection, payload: bytes) -> None:
        """Handle a binary WebSocket message."""
        client.session["_last_active"] = time.time()
        # Log binary message
        await self._log_websocket(client.id, 'in', event='binary',
                                  data=payload, opcode=0x2, size=len(payload))
        for handler in self._binary_handlers:
            try:
                await handler(client, payload)
            except Exception:
                self._logger.exception("Binary handler error")

    # ------------------------------------------------------------------
    # UDP request handling (upgraded with token-based rate limiting)
    # ------------------------------------------------------------------
    async def _handle_udp_datagram(self, data: bytes, addr: Tuple[str, int], client_ip: str, client_port: int) -> None:
        """Process a UDP datagram with optional client_id token for rate limiting."""
        # 1. Enforce max payload size
        if len(data) > self.max_payload_size:
            self._logger.debug(f"UDP datagram too large ({len(data)} bytes) from {addr}, dropping")
            return

        # 2. Parse JSON
        try:
            msg = json.loads(data)
            path = msg.get("path")
            if not isinstance(path, str):
                raise ValueError("Missing or invalid 'path'")
            body = msg.get("data", {})
            client_token = msg.get("client_id")  # optional
        except (json.JSONDecodeError, ValueError) as e:
            self._logger.debug(f"Invalid UDP message from {addr}: {e}")
            return

        # 3. Rate limiting: key = client_token if present, else IP:port
        rate_key = client_token if isinstance(client_token, str) else f"{client_ip}:{client_port}"
        if self.rate_limit:
            now = time.time()
            times = self._client_msg_times[rate_key]
            # Remove messages older than 1 second
            while times and times[0] < now - 1:
                times.popleft()
            if len(times) >= self.rate_limit:
                self._logger.debug(f"UDP rate limit exceeded for {rate_key}")
                return  # silently drop
            times.append(now)

        # 4. Build Request object
        req = Request(
            method="UDP",
            path=path,
            headers={},
            body=json.dumps(body).encode(),
            client_ip=client_ip
        )
        req.start_time = time.time()
        req.request_id = str(uuid.uuid4())
        self.request_id_ctx.set(req.request_id)

        # 5. Run HTTP middleware
        for mw in self._http_middleware:
            try:
                resp = await mw(req)
                if resp is not None:
                    await self._send_udp_response(addr, resp, req)
                    return
            except Exception as e:
                self._logger.exception("UDP middleware error")
                await self._send_udp_error(addr, 500, "Middleware error", req)
                return

        # 6. Find matching route
        route_path = path.split("?", 1)[0]
        for pattern, handler, methods in self._http_routes:
            match = pattern.match(route_path)
            if match:
                if "UDP" not in methods:
                    await self._send_udp_error(addr, 405, f"Method UDP not allowed", req)
                    return
                kwargs = match.groupdict()
                try:
                    result = await handler(req, **kwargs)
                    await self._send_udp_response(addr, result, req)
                except HTTPError as e:
                    await self._send_udp_error(addr, e.status_code, e.message, req)
                except Exception as e:
                    self._logger.exception(f"UDP handler error for {path}")
                    await self._send_udp_error(addr, 500, "Internal Server Error", req)
                break
        else:
            await self._send_udp_error(addr, 404, f"Not Found: {route_path}", req)

    async def _send_udp_response(self, addr: Tuple[str, int], result: Any, req: Request) -> None:
        """Convert handler result to a UDP datagram and send back."""
        if isinstance(result, dict) and "_redirect" in result:
            response_data = {"status": result.get("_status", 302), "location": result["_redirect"]}
            payload = json.dumps(response_data).encode()
            self._udp_transport.sendto(payload, addr)
            await self._log_http(req, result.get("_status", 302))
            return

        if isinstance(result, dict) and "_json" in result:
            response_data = result["_json"]
            status = result.get("_status", 200)
            payload = json.dumps(response_data).encode()
            self._udp_transport.sendto(payload, addr)
            await self._log_http(req, status)
            return

        if isinstance(result, tuple) and len(result) == 2:
            body, _ = result
        elif isinstance(result, str):
            body = result
        elif isinstance(result, dict):
            body = json.dumps(result)
        elif isinstance(result, bytes):
            body = result
        else:
            body = str(result)

        if isinstance(body, str):
            payload = body.encode()
        else:
            payload = body

        self._udp_transport.sendto(payload, addr)
        await self._log_http(req, 200)

    async def _send_udp_error(self, addr: Tuple[str, int], code: int, message: str, req: Request) -> None:
        error_payload = json.dumps({"error": code, "message": message}).encode()
        try:
            self._udp_transport.sendto(error_payload, addr)
        except:
            pass
        await self._log_http(req, code)

    # ------------------------------------------------------------------
    # HTTP response helpers
    # ------------------------------------------------------------------
    async def _send_http_response(self, writer: asyncio.StreamWriter, result: Any, req: Request) -> None:
        """Convert handler result to HTTP response and send."""
        # Handle redirect
        if isinstance(result, dict) and "_redirect" in result:
            location = result["_redirect"]
            status = result.get("_status", 302)
            response = (
                f"HTTP/1.1 {status} Found\r\n"
                f"Location: {location}\r\n"
                f"Content-Length: 0\r\n"
                f"Connection: close\r\n"
                f"\r\n"
            )
            writer.write(response.encode())
            await writer.drain()
            # Log HTTP response
            await self._log_http(req, status)
            return

        # Handle JSON response
        if isinstance(result, dict) and "_json" in result:
            body = json.dumps(result["_json"])
            content_type = "application/json"
            status = result.get("_status", 200)
            # Add CORS headers if enabled
            headers = []
            if self.cors_allowed_origins:
                origin = getattr(req, "_cors_origin", req.headers.get("origin", "*"))
                if origin in self.cors_allowed_origins or "*" in self.cors_allowed_origins:
                    headers.append(f"Access-Control-Allow-Origin: {origin}")
                    if self.cors_allow_credentials:
                        headers.append("Access-Control-Allow-Credentials: true")
            response = http_response(status, content_type, body)
            # Prepend CORS headers if any
            if headers:
                response_lines = response.split(b"\r\n")
                # Insert after status line
                for i, h in enumerate(headers):
                    response_lines.insert(1 + i, h.encode())
                response = b"\r\n".join(response_lines)
            writer.write(response)
            await writer.drain()
            # Log HTTP response
            await self._log_http(req, status)
            return

        # Handle FileResponse (streaming with custom headers)
        if isinstance(result, FileResponse):
            # Build base headers
            base_headers = [
                f"HTTP/1.1 200 OK",
                f"Content-Type: {result.content_type}",
                f"Content-Length: {result.size}",
                "Connection: close",
            ]
            # Merge custom headers (may override base ones)
            headers = {}
            for h in base_headers:
                if h.startswith("HTTP/"):
                    continue
                key, val = h.split(":", 1)
                headers[key.strip()] = val.strip()
            headers.update(result.headers)
            # Add CORS if enabled
            if self.cors_allowed_origins:
                origin = getattr(req, "_cors_origin", req.headers.get("origin", "*"))
                if origin in self.cors_allowed_origins or "*" in self.cors_allowed_origins:
                    headers["Access-Control-Allow-Origin"] = origin
                    if self.cors_allow_credentials:
                        headers["Access-Control-Allow-Credentials"] = "true"
            # Build response lines
            response_lines = [f"HTTP/1.1 200 OK"]
            for k, v in headers.items():
                response_lines.append(f"{k}: {v}")
            response_lines.append("\r\n")
            writer.write("\r\n".join(response_lines).encode())
            await writer.drain()
            async for chunk in result:
                writer.write(chunk)
                await writer.drain()
            # Log HTTP response
            await self._log_http(req, 200)
            return

        # Handle StreamingResponse (chunked)
        if isinstance(result, StreamingResponse):
            headers = [
                f"HTTP/1.1 200 OK",
                f"Content-Type: {result.content_type}",
                "Transfer-Encoding: chunked",
                "Connection: close",
            ]
            if self.cors_allowed_origins:
                origin = getattr(req, "_cors_origin", req.headers.get("origin", "*"))
                if origin in self.cors_allowed_origins or "*" in self.cors_allowed_origins:
                    headers.append(f"Access-Control-Allow-Origin: {origin}")
                    if self.cors_allow_credentials:
                        headers.append("Access-Control-Allow-Credentials: true")
            headers.append("\r\n")
            writer.write("\r\n".join(headers).encode())
            await writer.drain()
            async for chunk in result.generator:
                if chunk:
                    size = f"{len(chunk):X}\r\n".encode()
                    writer.write(size)
                    writer.write(chunk)
                    writer.write(b"\r\n")
                    await writer.drain()
            writer.write(b"0\r\n\r\n")
            await writer.drain()
            # Log HTTP response
            await self._log_http(req, 200)
            return

        # Handle tuple (body, content_type)
        if isinstance(result, tuple) and len(result) == 2:
            body, content_type = result
        elif isinstance(result, str):
            body = result
            content_type = "text/html"
        elif isinstance(result, dict):
            body = json.dumps(result)
            content_type = "application/json"
        elif isinstance(result, bytes):
            body = result
            content_type = "application/octet-stream"
        else:
            body = str(result)
            content_type = "text/plain"

        response = http_response(200, content_type, body)
        # Add CORS if needed
        if self.cors_allowed_origins:
            origin = getattr(req, "_cors_origin", req.headers.get("origin", "*"))
            if origin in self.cors_allowed_origins or "*" in self.cors_allowed_origins:
                # Insert Access-Control header
                response_lines = response.split(b"\r\n")
                cors_line = f"Access-Control-Allow-Origin: {origin}".encode()
                response_lines.insert(1, cors_line)
                if self.cors_allow_credentials:
                    response_lines.insert(2, b"Access-Control-Allow-Credentials: true")
                response = b"\r\n".join(response_lines)
        writer.write(response)
        await writer.drain()
        # Log HTTP response
        await self._log_http(req, 200)

    async def _send_http_error(self, writer: asyncio.StreamWriter, code: int, message: str, req: Request) -> None:
        """Send an HTTP error response and log it."""
        response = http_response(code, "text/plain", message)
        try:
            writer.write(response)
            await writer.drain()
        except:
            pass
        finally:
            writer.close()
            await writer.wait_closed()
            # Log HTTP error with the provided request object
            if req:
                await self._log_http(req, code)


# ----------------------------------------------------------------------
# send_file helper (enhanced)
# ----------------------------------------------------------------------
def send_file(
    filepath: str,
    base_dir: str = ".",
    content_type: Optional[str] = None,
    headers: Optional[Dict[str, str]] = None,
    as_attachment: bool = False,
    attachment_filename: Optional[str] = None,
    cache_control: Optional[str] = None,
    last_modified: Optional[float] = None,
) -> FileResponse:
    """
    Return a FileResponse that streams a file with advanced options.

    :param filepath: Relative path to the file (within base_dir or absolute if base_dir="")
    :param base_dir: Base directory (default: current working directory). Set to "" for absolute paths.
    :param content_type: Override MIME type.
    :param headers: Additional HTTP headers as dict.
    :param as_attachment: If True, force download with Content-Disposition.
    :param attachment_filename: Filename used in Content-Disposition (defaults to base name of file).
    :param cache_control: Cache-Control header value (e.g., "max-age=3600, public").
    :param last_modified: Override Last-Modified timestamp (Unix timestamp).
    :raises FileNotFoundError: If file does not exist or is a directory.
    """
    if base_dir:
        full_path = os.path.join(base_dir, filepath)
    else:
        full_path = filepath
    if not os.path.exists(full_path) or os.path.isdir(full_path):
        raise FileNotFoundError(f"File not found: {full_path}")
    return FileResponse(
        filepath=full_path,
        content_type=content_type,
        headers=headers,
        as_attachment=as_attachment,
        attachment_filename=attachment_filename,
        cache_control=cache_control,
        last_modified=last_modified,
    )


# ----------------------------------------------------------------------
# Template rendering helper
# ----------------------------------------------------------------------
def render_template(template_name: str, context: Optional[Dict[str, Any]] = None, template_dir: str = "templates") -> str:
    """
    Render an HTML template with variable substitution.
    Supports nested keys using dot notation, e.g. {{ user.name }}.

    Looks for the template file in `template_dir` (relative to current working directory).
    Variables in the template should be written as `{{ variable }}` or `{{ variable.nested }}`.
    Returns the rendered HTML as a string.
    """
    if context is None:
        context = {}

    template_path = os.path.join(template_dir, template_name)
    try:
        with open(template_path, 'r', encoding='utf-8') as f:
            content = f.read()
    except FileNotFoundError:
        raise FileNotFoundError(f"Template not found: {template_path}")

    def replace(match):
        expr = match.group(1).strip()
        parts = expr.split('.')
        value = context
        for part in parts:
            if isinstance(value, dict) and part in value:
                value = value[part]
            else:
                # path not found, return the original placeholder unchanged
                return match.group(0)
        return str(value)

    # Pattern matches {{ variable }} with optional dots, allowing spaces inside braces
    pattern = r'{{\s*([^}\s]+(?:\.[^}\s]+)*)\s*}}'
    return re.sub(pattern, replace, content)


# ----------------------------------------------------------------------
# CLI entry point (updated with --protocol AUTO)
# ----------------------------------------------------------------------
if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Lynk server runner")
    parser.add_argument("app", help="Application module in format 'module:app'")
    parser.add_argument("--host", default="0.0.0.0", help="Host to bind")
    parser.add_argument("--port", type=int, default=8765, help="Port to bind")
    parser.add_argument("--protocol", default="TCP", choices=["TCP", "UDP", "AUTO"], help="Protocol to use")
    parser.add_argument("--debug", action="store_true", help="Enable debug mode")
    args = parser.parse_args()

    # Import the app dynamically
    module_path, app_name = args.app.split(":")
    import importlib
    module = importlib.import_module(module_path)
    app = getattr(module, app_name)

    # Configure logging
    logging.basicConfig(level=logging.DEBUG if args.debug else logging.INFO)

    # Set host/port/protocol if app doesn't have them
    if hasattr(app, 'host'):
        app.host = args.host
    if hasattr(app, 'port'):
        app.port = args.port
    if hasattr(app, 'protocol'):
        app.protocol = args.protocol   # override if present
    if hasattr(app, 'debug'):
        app.debug = args.debug

    app.run()