# SPDX-FileCopyrightText: 2026 Bart Broere
# SPDX-License-Identifier: GPL-3.0-or-later
"""Bolt protocol server (raw TCP and WebSocket on the same port).

Supports Bolt 4.4 and 5.0-5.4, which covers the neo4j-javascript-driver used
by neo4j-browser as well as current Python/Java drivers and cypher-shell.
"""

import asyncio
import itertools
import logging
import struct

from websockets.frames import CloseCode, Opcode
from websockets.protocol import State
from websockets.server import ServerProtocol

from . import packstream
from .backend import BoltFailure, Config, Session
from .discovery import discovery_document
from .packstream import Structure

log = logging.getLogger("peo.bolt")

MAGIC = b"\x60\x60\xB0\x17"
SUPPORTED = [(5, 4), (5, 3), (5, 2), (5, 1), (5, 0), (4, 4)]
SERVER_AGENT = "Neo4j/5.26.0"

HELLO, GOODBYE, RESET, RUN, BEGIN, COMMIT, ROLLBACK = 0x01, 0x02, 0x0F, 0x10, 0x11, 0x12, 0x13
DISCARD, PULL, TELEMETRY, ROUTE, LOGON, LOGOFF = 0x2F, 0x3F, 0x54, 0x66, 0x6A, 0x6B
SUCCESS, RECORD, IGNORED, FAILURE = 0x70, 0x71, 0x7E, 0x7F

_connection_ids = itertools.count(1)


class ConnectionClosed(Exception):
    pass


# -- transports -----------------------------------------------------------------


class RawTransport:
    def __init__(self, reader, writer, prefix=b""):
        self.reader, self.writer, self.buffer = reader, writer, bytearray(prefix)

    async def read(self, n):
        while len(self.buffer) < n:
            chunk = await self.reader.read(65536)
            if not chunk:
                raise ConnectionClosed()
            self.buffer += chunk
        data = bytes(self.buffer[:n])
        del self.buffer[:n]
        return data

    async def write(self, data):
        self.writer.write(data)
        await self.writer.drain()

    async def close(self):
        self.writer.close()


class WebSocketTransport(RawTransport):
    """Bolt over WebSocket: Bolt bytes are carried in binary frames.

    Uses the Sans-I/O protocol of the websockets library, as this port also speaks raw Bolt.
    """

    async def handshake(self):
        # Read the request head here rather than in websockets, which only accepts GET:
        # neo4j-browser also sends plain HTTP requests (discovery, CORS preflight) to this port.
        request = bytearray(self.buffer)
        self.buffer.clear()
        while b"\r\n\r\n" not in request:
            chunk = await self.reader.read(4096)
            if not chunk or len(request) > 65536:
                raise ConnectionClosed()
            request += chunk
        head = bytes(request).partition(b"\r\n\r\n")[0]
        headers = {}
        for line in head.decode("latin-1").split("\r\n")[1:]:
            k, _, v = line.partition(":")
            headers[k.strip().lower()] = v.strip()
        method = head.split(b" ", 1)[0]
        if method != b"GET" or "websocket" not in headers.get("upgrade", "").lower():
            await self._plain_http(method, headers)
            raise ConnectionClosed()

        self.protocol = ServerProtocol(max_size=None)
        self.protocol.receive_data(bytes(request))
        events = self.protocol.events_received()
        if not events:
            raise ConnectionClosed()
        response = self.protocol.accept(events[0])
        self.protocol.send_response(response)
        await self._flush()
        if response.status_code != 101:
            raise ConnectionClosed()

    async def _plain_http(self, method, headers):
        """Answer non-WebSocket requests, e.g. neo4j-browser probing for discovery data."""
        cors = (b"Access-Control-Allow-Origin: *\r\nAccess-Control-Allow-Headers: *\r\n"
                b"Access-Control-Allow-Methods: GET, HEAD, OPTIONS\r\nConnection: close\r\n")
        if method == b"OPTIONS":
            response = b"HTTP/1.1 204 No Content\r\n" + cors + b"\r\n"
        elif method in (b"GET", b"HEAD"):
            body = discovery_document(headers, self.writer.get_extra_info("sockname")[1])
            response = (b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\n" + cors
                        + b"Content-Length: %d\r\n\r\n" % len(body) + (body if method == b"GET" else b""))
        else:
            response = b"HTTP/1.1 405 Method Not Allowed\r\n" + cors + b"Content-Length: 0\r\n\r\n"
        self.writer.write(response)
        await self.writer.drain()

    async def _flush(self):
        for data in self.protocol.data_to_send():
            if data:
                self.writer.write(data)
        await self.writer.drain()

    async def read(self, n):
        while len(self.buffer) < n:
            for frame in self.protocol.events_received():
                if frame.opcode in (Opcode.CONT, Opcode.TEXT, Opcode.BINARY):
                    self.buffer += frame.data
            await self._flush()  # pongs and the reply to a close frame
            if len(self.buffer) >= n:
                break
            if self.protocol.state is not State.OPEN:
                raise ConnectionClosed()
            chunk = await self.reader.read(65536)
            if not chunk:
                raise ConnectionClosed()
            self.protocol.receive_data(chunk)
        data = bytes(self.buffer[:n])
        del self.buffer[:n]
        return data

    async def write(self, data):
        self.protocol.send_binary(data)
        await self._flush()

    async def close(self):
        if self.protocol.state is State.OPEN:
            self.protocol.send_close(CloseCode.NORMAL_CLOSURE)
            await self._flush()
        self.writer.close()


# -- protocol -----------------------------------------------------------------------


def negotiate(proposals: bytes):
    for i in range(4):
        _, span, minor, major = proposals[i * 4:i * 4 + 4]
        for version in SUPPORTED:
            if version[0] == major and minor - span <= version[1] <= minor:
                return version
    return None


class BoltConnection:
    def __init__(self, transport, config: Config):
        self.t = transport
        self.config = config
        self.id = f"bolt-{next(_connection_ids)}"
        self.version = None
        self.packer = None
        self.session: Session | None = None
        self.failed = False
        self.routing = {}
        self.results = {}  # qid -> [Result, next row index]
        self.last_qid = -1
        self.qids = itertools.count()

    async def serve(self):
        proposals = await self.t.read(16)
        self.version = negotiate(proposals)
        if self.version is None:
            log.info("%s: no supported Bolt version in %s", self.id, proposals.hex())
            await self.t.write(b"\x00\x00\x00\x00")
            return
        await self.t.write(bytes([0, 0, self.version[1], self.version[0]]))
        self.packer = packstream.Packer(element_ids=self.version >= (5, 0))
        log.debug("%s: negotiated Bolt %d.%d", self.id, *self.version)
        try:
            while True:
                message = await self.read_message()
                if message is None:
                    continue  # NOOP chunk
                if not await self.handle(message):
                    break
        finally:
            if self.session:
                await self.session.close()

    async def read_message(self):
        data = bytearray()
        while True:
            size = struct.unpack(">H", await self.t.read(2))[0]
            if size == 0:
                if not data:
                    return None
                return packstream.unpack(bytes(data))
            data += await self.t.read(size)

    async def send(self, tag, *fields):
        data = self.packer.pack(Structure(tag, list(fields)))
        out = bytearray()
        for i in range(0, len(data), 0xFFFF):
            chunk = data[i:i + 0xFFFF]
            out += struct.pack(">H", len(chunk)) + chunk
        out += b"\x00\x00"
        await self.t.write(bytes(out))

    async def success(self, metadata=None):
        await self.send(SUCCESS, metadata or {})

    async def failure(self, code, message):
        log.info("%s: %s: %s", self.id, code, message)
        self.failed = True
        await self.send(FAILURE, {"code": code, "message": message})

    async def handle(self, msg) -> bool:
        tag, f = msg.tag, msg.fields
        if tag == GOODBYE:
            return False
        if tag == RESET:
            self.failed = False
            self.results.clear()
            if self.session:
                await self.session.reset()
            await self.success()
            return True
        if self.failed:
            await self.send(IGNORED)
            return True
        try:
            await self.dispatch(tag, f)
        except BoltFailure as e:
            self.results.clear()
            await self.failure(e.code, e.message)
        except Exception as e:  # noqa: BLE001 - report anything to the client
            log.exception("%s: internal error", self.id)
            self.results.clear()
            await self.failure("Neo.DatabaseError.General.UnknownError", f"peo internal error: {e}")
        return True

    async def dispatch(self, tag, f):
        if tag == HELLO:
            extra = f[0] if f else {}
            self.routing = extra.get("routing") or {}
            if self.version < (5, 1):
                await self.logon(extra)
            await self.success({"server": SERVER_AGENT, "connection_id": self.id, "hints": {}})
        elif tag == LOGON:
            await self.logon(f[0])
            await self.success()
        elif tag == LOGOFF:
            if self.session:
                await self.session.close()
            self.session = None
            await self.success()
        elif tag == TELEMETRY:
            await self.success()
        elif tag == ROUTE:
            await self.route(f)
        elif tag == RUN:
            await self.run(f)
        elif tag in (PULL, DISCARD):
            await self.pull(f[0] if f else {}, discard=tag == DISCARD)
        elif tag == BEGIN:
            extra = f[0] if f else {}
            await self.require_session().begin(extra.get("db"))
            self.results.clear()
            await self.success()
        elif tag == COMMIT:
            await self.require_session().commit()
            await self.success({"bookmark": "peo:bookmark"})
        elif tag == ROLLBACK:
            await self.require_session().rollback()
            await self.success()
        else:
            raise BoltFailure("Neo.ClientError.Request.Invalid", f"Unsupported Bolt message 0x{tag:02X}")

    def require_session(self) -> Session:
        if self.session is None:
            raise BoltFailure("Neo.ClientError.Security.Unauthorized", "Not authenticated")
        return self.session

    async def logon(self, auth):
        if auth.get("scheme") != "basic":
            raise BoltFailure("Neo.ClientError.Security.Unauthorized",
                              "peo requires basic authentication with a PostgreSQL user name and password.")
        session = Session(self.config, auth.get("principal", ""))
        await session.connect(auth.get("credentials", ""))
        self.session = session

    async def route(self, f):
        extra = f[2] if len(f) > 2 else None
        db = extra.get("db") if isinstance(extra, dict) else extra
        db = await self.require_session().resolve_db(db)
        address = self.routing.get("address") or self.config.advertised_bolt
        await self.success({"rt": {"ttl": 300, "db": db, "servers": [
            {"addresses": [address], "role": "WRITE"},
            {"addresses": [address], "role": "READ"},
            {"addresses": [address], "role": "ROUTE"},
        ]}})

    async def run(self, f):
        query, params = f[0], f[1] if len(f) > 1 else {}
        extra = f[2] if len(f) > 2 else {}
        session = self.require_session()
        if not session.in_transaction:
            self.results.clear()
        result = await session.run(query, params, extra.get("db"))
        qid = next(self.qids)
        self.results[qid] = [result, 0]
        self.last_qid = qid
        metadata = {"fields": result.fields, "t_first": result.t_first}
        if session.in_transaction:
            metadata["qid"] = qid
        await self.success(metadata)

    async def pull(self, extra, discard):
        qid = extra.get("qid", -1)
        qid = self.last_qid if qid == -1 else qid
        if qid not in self.results:
            raise BoltFailure("Neo.ClientError.Request.Invalid", "No result available to pull from")
        result, index = self.results[qid]
        n = extra.get("n", -1)
        end = len(result.rows) if n == -1 or discard else min(index + n, len(result.rows))
        if not discard:
            for row in result.rows[index:end]:
                await self.send(RECORD, row)
        self.results[qid][1] = end
        if end < len(result.rows):
            await self.success({"has_more": True})
            return
        del self.results[qid]
        metadata = {"type": result.query_type, "t_last": 0, "db": result.db}
        if not self.session.in_transaction:
            metadata["bookmark"] = "peo:bookmark"
        await self.success(metadata)


async def handle_client(reader, writer, config: Config):
    peer = writer.get_extra_info("peername")
    try:
        first = await reader.readexactly(4)
        log.debug("connection from %s starts with %r", peer, first)
        if first == MAGIC:
            transport = RawTransport(reader, writer)
        elif first in (b"GET ", b"HEAD", b"OPTI", b"POST"):
            transport = WebSocketTransport(reader, writer, first)
            await transport.handshake()
            if await transport.read(4) != MAGIC:
                return
        else:
            return
        await BoltConnection(transport, config).serve()
        await transport.close()
    except (ConnectionClosed, asyncio.IncompleteReadError, ConnectionResetError, BrokenPipeError):
        pass
    except Exception:  # noqa: BLE001
        log.exception("connection from %s failed", peer)
    finally:
        writer.close()