"""
Rembus messages, protocol constants and utility function related
to the transport layer.
"""

import base64
import os
import logging
from typing import Any, List
import json
import time
import cbor2
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa, ec
from cryptography.hazmat.primitives.asymmetric.types import PrivateKeyTypes
from cryptography.hazmat.backends import default_backend
from narwhals.typing import IntoFrame
import polars as pl
import pandas as pd
import pyarrow as pa
from pydantic import BaseModel, PrivateAttr
import rembus.settings as rs

__all__ = [
    "CBOR",
    "JSON",
    "RembusMsg",
    "Ack2Msg",
    "AckMsg",
    "AdminMsg",
    "AttestationMsg",
    "IdentityMsg",
    "PubSubMsg",
    "RegisterMsg",
    "ResMsg",
    "RpcReqMsg",
    "UnregisterMsg",
]

logger = logging.getLogger(__name__)

WS_FRAME_MAXSIZE = 60 * 1024 * 1024

SIG_RSA = 0x1
SIG_ECDSA = 0x2

CBOR = 0
"""CBOR encoding"""

JSON = 1
"""JSON encoding"""

Now = 0.0
LastReceived = float("inf")

QOS0: int = 0x00
QOS1: int = 0x10
QOS2: int = 0x30
SLOT_FLAG: int = 0x40

MSGID_SZ = 8

TYPE_IDENTITY = 0
TYPE_PUB = 1
TYPE_RPC = 2
TYPE_ADMIN = 3
TYPE_RESPONSE = 4
TYPE_ACK = 5
TYPE_ACK2 = 6
TYPE_UNREGISTER = 9
TYPE_REGISTER = 10
TYPE_ATTESTATION = 11

STS_OK = 0
STS_ERROR = 0x0A
STS_CHALLENGE = 0x0B  # 11
STS_IDENTIFICATION_ERROR = 0x14  # 20
STS_METHOD_EXCEPTION = 0x28  # 40
STS_METHOD_ARGS_ERROR = 0x29  # 41
STS_METHOD_NOT_FOUND = 0x2A  # 42
STS_METHOD_UNAVAILABLE = 0x2B  # 43
STS_METHOD_LOOPBACK = 0x2C  # 44
STS_TARGET_NOT_FOUND = 0x2D  # 45
STS_TARGET_DOWN = 0x2E  # 46
STS_UNKNOWN_ADMIN_CMD = 0x2F  # 47
STS_NAME_ALREADY_TAKEN = 0x3C  # 60

DATAFRAME_TAG = 80

retcode = {
    STS_OK: "OK",
    STS_ERROR: "internal error",
    STS_IDENTIFICATION_ERROR: "IDENTIFICATION_ERROR",
    STS_METHOD_EXCEPTION: "METHOD_EXCEPTION",
    STS_METHOD_ARGS_ERROR: "METHOD_ARGS_ERROR",
    STS_METHOD_NOT_FOUND: "METHOD_NOT_FOUND",
    STS_METHOD_UNAVAILABLE: "METHOD_UNAVAILABLE",
    STS_METHOD_LOOPBACK: "METHOD_LOOPBACK",
    STS_TARGET_NOT_FOUND: "TARGET_NOT_FOUND",
    STS_TARGET_DOWN: "TARGET_DOWN",
    STS_UNKNOWN_ADMIN_CMD: "UNKNOWN_ADMIN_CMD",
    STS_NAME_ALREADY_TAKEN: "NAME_ALREADY_TAKEN",
}

BROKER_CONFIG = "__config__"
CID = "cid"
COMMAND = "cmd"
MSG_FROM = "msg_from"
ADD_INTEREST = "subscribe"
REMOVE_INTEREST = "unsubscribe"
ADD_IMPL = "expose"
REMOVE_IMPL = "unexpose"
REACTIVE_CMD = "reactive"
PRIVATE_TOPIC = "private_topic"
PUBLIC_TOPIC = "public_topic"
AUTHORIZE = "authorize"
UNAUTHORIZE = "unauthorize"
STATUS = "status"


def timestamp():
    return time.time_ns()


def msgid():
    """Return an array of MSGID_SZ random bytes."""
    return int.from_bytes(os.urandom(MSGID_SZ))


def msgid_slot(slot: int):
    """Return MSGID_SZ-4 random bytes and a 4 bytes slot embedded value."""
    mid = int.from_bytes(os.urandom(MSGID_SZ))
    return mid & 0xFFFFFFFF00000000 | (slot & 0xFFFFFFFF)


def bytes2id(byte_data: bytearray) -> int:
    """
    Converts a MSGID_SZ-byte bytearray to a UInt128 integer.

    Args:
        byte_data: A bytearray of MSGID_SZ bytes.

    Returns:
        A Python integer representing the UInt128 value.
    """
    if len(byte_data) != MSGID_SZ:
        raise ValueError(f"bytearray must be exactly {MSGID_SZ} bytes long")

    # Convert bytes to integer (assuming little-endian, adjust if big-endian)
    return from_bytes(byte_data)


class SendDataAtRest:
    def __init__(self, twin):
        self.twin = twin


class RembusException(Exception):
    """Base class for all Rembus exceptions."""


class RembusTimeout(RembusException):
    """Raised when a Rembus message that expects a response times out."""

    def __str__(self):
        return "request timeout"


class RembusConnectionClosed(RembusException):
    """Raised when a Rembus connection is closed unexpectedly."""

    def __str__(self):
        return "connection down"


class RembusError(RembusException):
    """Raised when a Rembus message returns a generic error."""

    def __init__(self, status_code: int, msg: str | None = None):
        self.status = status_code
        self.message = msg

    def __str__(self):
        if self.message:
            return f"{retcode[self.status]}:{self.message}"
        else:
            return f"{retcode[self.status]}"


def to_bytes(val: int) -> bytes:
    """Convert an int to MSGID_SZ bytes"""
    return val.to_bytes(MSGID_SZ, "little")


def from_bytes(bval: bytes | bytearray) -> int:
    """Convert MSGID_SZ bytes to an int"""
    return int.from_bytes(bval, "little")


class RembusMsg(BaseModel):
    """Rembus message"""

    _twin: Any = PrivateAttr(default=None)

    @property
    def twin(self):
        """Return the Twin that owns the message"""
        return self._twin

    @twin.setter
    def twin(self, rb):
        self._twin = rb

    def to_payload(self, enc: int) -> bytes | str:
        """
        Encode the message into a serializable format.

        Parameters
        ----------
        enc : int
            The encoding format. Supported values:

            * :obj:`CBOR` : return a CBOR-encoded byte string
            * :obj:`JSON` : return a JSON-encoded string

        Returns
        -------
        bytes | str
            The serialized payload. Returns ``bytes`` for CBOR encoding
            and ``str`` for JSON encoding.

        Notes
        -----
        This is an abstract method and must be implemented in subclasses.
        It is available in all subclasses of `RembusMsg` such as
        :class:`RpcReqMsg` and :class:`ResMsg`.
        """
        raise NotImplementedError("Subclasses must implement to_payload")


class RpcReqMsg(RembusMsg):
    """RPC request message."""

    id: int
    topic: str
    data: Any = None
    target: str | None = None

    def to_payload(self, enc: int) -> bytes | str:
        """Serialize the RpcRequest message."""
        if enc == CBOR:
            return cbor2.dumps(
                [
                    TYPE_RPC,
                    to_bytes(self.id),
                    self.topic,
                    self.target,
                    self.data,
                ]
            )

        return json.dumps(
            {
                "jsonrpc": "2.0",
                "id": self.id,
                "method": self.topic,
                "params": self.data,
            }
        )


class ResMsg(RembusMsg):
    """Response message."""

    id: int
    status: int
    data: Any = None
    _reqdata: Any = PrivateAttr()

    def to_payload(self, enc: int) -> bytes | str:
        """Serialize the Response message."""
        if enc == CBOR:
            return cbor2.dumps(
                [
                    TYPE_RESPONSE,
                    to_bytes(self.id),
                    self.status,
                    df2tag(self.data),
                ]
            )

        if self.status == STS_OK or self.status == STS_CHALLENGE:
            restype = "result"
        else:
            restype = "error"

        return json.dumps(
            {
                "jsonrpc": "2.0",
                "id": self.id,
                restype: {
                    "type": TYPE_RESPONSE,
                    "sts": self.status,
                    "data": self.data,
                },
            }
        )


class PubSubMsg(RembusMsg):
    """Pub/Sub message."""

    id: int | None = None
    topic: str
    data: Any = None
    flags: int = QOS0  # type: ignore[valid-type]
    slot: int | None = None
    recvts: int | None = None
    table: str | None = None
    regex: str | None = None

    def to_payload(self, enc: int) -> bytes | str:
        """Serialize the Pub/Sub message."""

        if self.id is not None:
            if self.slot is not None:
                self.flags |= SLOT_FLAG

            if enc == CBOR:
                return cbor2.dumps(
                    [
                        TYPE_PUB | self.flags,
                        to_bytes(self.id),
                        self.topic,
                        df2tag(self.data),
                    ]
                )

            return json.dumps(
                {
                    "jsonrpc": "2.0",
                    "id": self.id,
                    "method": self.topic,
                    "params": {"type": self.flags, "data": self.data},
                }
            )
        else:
            if self.slot is None:
                if enc == CBOR:
                    return cbor2.dumps(
                        [TYPE_PUB, self.topic, df2tag(self.data)]
                    )

                return json.dumps(
                    {
                        "jsonrpc": "2.0",
                        "method": self.topic,
                        "params": self.data,
                    }
                )
            else:
                self.flags |= SLOT_FLAG
                mid = to_bytes(self.slot)
                if enc == CBOR:
                    return cbor2.dumps(
                        [
                            TYPE_PUB | self.flags,
                            mid,
                            self.topic,
                            df2tag(self.data),
                        ]
                    )

                return json.dumps(
                    {
                        "jsonrpc": "2.0",
                        "id": self.slot,
                        "method": self.topic,
                        "params": self.data,
                    }
                )


class AckMsg(RembusMsg):
    """Ack message used for Pub/Sub QOS1 and QOS2."""

    id: int

    def to_payload(self, enc: int) -> bytes | str:
        """Serialize the Ack message."""
        if enc == CBOR:
            return cbor2.dumps([TYPE_ACK, to_bytes(self.id)])

        return json.dumps(
            {"jsonrpc": "2.0", "id": self.id, "result": {"type": TYPE_ACK}}
        )


class Ack2Msg(RembusMsg):
    """Ack2 message used for QOS2"""

    id: int

    def to_payload(self, enc: int) -> bytes | str:
        """Serialize the Ack2 message."""
        if enc == CBOR:
            return cbor2.dumps([TYPE_ACK2, to_bytes(self.id)])

        return json.dumps(
            {"jsonrpc": "2.0", "id": self.id, "result": {"type": TYPE_ACK2}}
        )


class AdminMsg(RembusMsg):
    """Admin message."""

    id: int
    topic: str
    data: Any = None

    def __str__(self):
        return f"AdminMsg:{self.topic}: {self.data}"

    def to_payload(self, enc: int) -> bytes | str:
        """Serialize the Admin message."""
        if enc == CBOR:
            return cbor2.dumps(
                [TYPE_ADMIN, to_bytes(self.id), self.topic, self.data]
            )

        return json.dumps(
            {
                "jsonrpc": "2.0",
                "id": self.id,
                "method": self.topic,
                "params": {"type": TYPE_ADMIN, "data": self.data},
            }
        )


class IdentityMsg(RembusMsg):
    """Identity message.

    This message is sent by the component to identify itself
    to the remote peer.
    """

    id: int
    cid: str

    def __str__(self):
        return f"IdentityMsg:{self.cid}"

    def to_payload(self, enc: int) -> bytes | str:
        """Serialize the Identity message."""
        if enc == CBOR:
            return cbor2.dumps([TYPE_IDENTITY, to_bytes(self.id), self.cid])

        return json.dumps(
            {
                "jsonrpc": "2.0",
                "id": self.id,
                "method": "__identity__",
                "params": {"type": TYPE_IDENTITY, "cid": self.cid},
            }
        )


class AttestationMsg(RembusMsg):
    """Attestation message."""

    id: int
    cid: str
    signature: bytes | str

    def __str__(self):
        return f"AttestationMsg:{self.cid}"

    def to_payload(self, enc: int) -> bytes | str:
        """Serialize the Attestation message."""
        if enc == CBOR:
            return cbor2.dumps(
                [TYPE_ATTESTATION, to_bytes(self.id), self.cid, self.signature]
            )

        return json.dumps(
            {
                "jsonrpc": "2.0",
                "id": self.id,
                "method": self.cid,
                "params": {
                    "type": TYPE_ATTESTATION,
                    "signature": self.signature,
                },
            }
        )


class RegisterMsg(RembusMsg):
    """Register message."""

    id: int
    cid: str
    pin: str
    pubkey: bytes | str
    type: int

    def to_payload(self, enc: int) -> bytes | str:
        """Serialize the Register message."""
        if enc == CBOR:
            return cbor2.dumps(
                [
                    TYPE_REGISTER,
                    to_bytes(self.id),
                    self.cid,
                    self.pin,
                    self.pubkey,
                    self.type,
                ]
            )

        return json.dumps(
            {
                "jsonrpc": "2.0",
                "id": self.id,
                "method": self.cid,
                "params": {
                    "type": TYPE_REGISTER,
                    "pin": self.pin,
                    "key_val": self.pubkey,
                    "key_type": self.type,
                },
            }
        )


class UnregisterMsg(RembusMsg):
    """Unregister message."""

    id: int

    def to_payload(self, enc: int) -> bytes | str:
        """Serialize the Unregister message."""
        if enc == CBOR:
            return cbor2.dumps([TYPE_UNREGISTER, to_bytes(self.id)])

        return json.dumps(
            {
                "jsonrpc": "2.0",
                "id": self.id,
                "method": "__unregister__",
                "params": {"type": TYPE_UNREGISTER},
            }
        )


def isregistered(router_id, rid: str):
    """Check if a component identified by `rid` is registered in the router."""
    try:
        rs.key_file(router_id, rid)
        return True
    except FileNotFoundError:
        return False


def tohex(bs: bytes):
    """Return a string with bytes as hex numbers with 0xNN format."""
    return " ".join(f"0x{x:02x}" for x in bs)


def jsonprc_request(pkt, msg_id, params) -> RembusMsg:
    """Parse a JSON_RPC request"""
    if isinstance(params, list):
        # Default to RPC request
        return RpcReqMsg(id=msg_id, topic=pkt["method"], data=params)
    else:
        msg_type = params.get("type")
        if msg_type in [QOS1, QOS2]:
            return PubSubMsg(
                id=msg_id,
                topic=pkt["method"],
                data=pkt.get("data"),
                flags=msg_type,
            )
        elif msg_type == TYPE_IDENTITY:
            return IdentityMsg(id=msg_id, cid=params["cid"])
        elif msg_type == TYPE_ADMIN:
            return AdminMsg(
                id=msg_id, topic=pkt["method"], data=params.get("data")
            )
        elif msg_type == TYPE_ATTESTATION:
            sig = params.get("signature")
            return AttestationMsg(id=msg_id, cid=pkt["method"], signature=sig)
        elif msg_type == TYPE_REGISTER:
            pubkey = params.get("key_val")
            if isinstance(pubkey, str):
                pubkey = base64.b64decode(pubkey)

            return RegisterMsg(
                id=msg_id,
                cid=pkt["method"],
                pin=params.get("pin"),
                pubkey=pubkey,
                type=params.get("key_type"),
            )
        elif msg_type == TYPE_UNREGISTER:
            return UnregisterMsg(id=msg_id)

        raise ValueError(f"{pkt}:invalid JSON-RPC request")


def jsonprc_response(pkt, msg_id, result) -> RembusMsg:
    """Parse a JSON_RPC success response"""

    msg_type = result.get("type")
    if msg_type == TYPE_RESPONSE or msg_type is None:
        status = result.get("sts", STS_OK)
        return ResMsg(id=msg_id, status=status, data=result.get("data"))
    elif msg_type == TYPE_ACK:
        return AckMsg(id=msg_id)
    elif msg_type == TYPE_ACK2:
        return Ack2Msg(id=msg_id)

    raise ValueError(f"{pkt}:invalid JSON-RPC response")


def jsonrpc_parse(payload) -> RembusMsg:
    """Get a Rembus message from a JSON-RPC string payload."""
    pkt = json.loads(payload)
    msg_id = pkt.get("id")
    if msg_id:
        # request-response message
        params = pkt.get("params")
        if params is not None:
            return jsonprc_request(pkt, msg_id, params)

        result = pkt.get("result")
        if result:
            return jsonprc_response(pkt, msg_id, result)

        err = pkt.get("error")
        if err:
            return jsonprc_response(pkt, msg_id, err)

    else:
        return PubSubMsg(topic=pkt["method"], data=pkt.get("data"))

    raise ValueError(f"{pkt}:JSON-RPC invalid payload")


def cbor_parse(pkt) -> RembusMsg:
    """Parse a CBOR packet into a Rembus message."""
    type_byte = pkt[0]
    mtype = type_byte & 0x0F
    flags = type_byte & 0xF0

    # Special cases first (because of flag-dependent structure)
    if mtype == TYPE_PUB:
        if flags == QOS0:
            return PubSubMsg(topic=pkt[1], data=pkt[2])

        slot = None
        mid = from_bytes(pkt[1])
        if flags & SLOT_FLAG:
            slot = mid & 0xFFFFFFFF

        return PubSubMsg(
            id=mid, slot=slot, topic=pkt[2], data=pkt[3], flags=flags
        )

    if mtype == TYPE_RESPONSE:
        return ResMsg(
            id=from_bytes(pkt[1]),
            status=pkt[2],
            data=pkt[3] if len(pkt) > 3 else None,
        )

    # Dispatch table for the rest
    table = {
        TYPE_RPC: lambda: RpcReqMsg(
            id=from_bytes(pkt[1]), topic=pkt[2], target=pkt[3], data=pkt[4]
        ),
        TYPE_ACK: lambda: AckMsg(id=from_bytes(pkt[1])),
        TYPE_ACK2: lambda: Ack2Msg(id=from_bytes(pkt[1])),
        TYPE_ADMIN: lambda: AdminMsg(
            id=from_bytes(pkt[1]), topic=pkt[2], data=pkt[3]
        ),
        TYPE_IDENTITY: lambda: IdentityMsg(id=from_bytes(pkt[1]), cid=pkt[2]),
        TYPE_ATTESTATION: lambda: AttestationMsg(
            id=from_bytes(pkt[1]), cid=pkt[2], signature=pkt[3]
        ),
        TYPE_REGISTER: lambda: RegisterMsg(
            id=from_bytes(pkt[1]),
            cid=pkt[2],
            pin=pkt[3],
            pubkey=pkt[4],
            type=pkt[5],
        ),
        TYPE_UNREGISTER: lambda: UnregisterMsg(id=from_bytes(pkt[1])),
    }

    try:
        return table[mtype]()
    except KeyError as e:
        raise ValueError(f"unknown message type: {mtype}") from e


def decode_dataframe(data: bytes) -> IntoFrame:
    """Decode a CBOR tagged value `data` to a pandas dataframe."""
    writer = pa.BufferOutputStream()
    writer.write(data)
    buf: pa.Buffer = writer.getvalue()
    with pa.ipc.open_stream(buf) as reader:
        table = reader.read_all()

    return pl.from_arrow(table)


def encode_dataframe(df: IntoFrame) -> cbor2.CBORTag:
    """Encode a pandas dataframe `df` to a CBOR tag value."""
    if isinstance(df, pl.DataFrame):
        table = df.to_arrow()
    else:
        table = pa.Table.from_pandas(df)

    sink = pa.BufferOutputStream()
    with pa.ipc.new_stream(sink, table.schema) as writer:
        writer.write(table)
    buf = sink.getvalue()
    stream = pa.input_stream(buf)
    return cbor2.CBORTag(DATAFRAME_TAG, stream.read())


def encode(msg: list[Any]) -> bytes:
    """Encode message `msg`."""
    return cbor2.dumps(msg)


def tag2df(data: Any) -> Any:
    """Loop over `data` items and decode tagged values to dataframes."""
    if isinstance(data, list):
        for idx, val in enumerate(data):
            if isinstance(val, cbor2.CBORTag) and val.tag == DATAFRAME_TAG:
                data[idx] = decode_dataframe(val.value)
    elif isinstance(data, cbor2.CBORTag) and data.tag == DATAFRAME_TAG:
        return decode_dataframe(data.value)
    return data


def df2tag(data: Any) -> Any:
    """Loop over `data` items and encode dataframes to tag values."""
    if isinstance(data, tuple):
        lst: List[Any] = []
        for idx, val in enumerate(data):
            if isinstance(val, pl.DataFrame) or isinstance(val, pd.DataFrame):
                lst.append(encode_dataframe(val))
            else:
                lst.append(val)
        return lst
    elif isinstance(data, list):
        for idx, val in enumerate(data):
            if isinstance(val, pl.DataFrame) or isinstance(val, pd.DataFrame):
                data[idx] = encode_dataframe(val)

    elif isinstance(data, pl.DataFrame) or isinstance(data, pd.DataFrame):
        data = encode_dataframe(data)
    return data


def rsa_private_key():
    """Generate a new RSA private key."""
    return rsa.generate_private_key(public_exponent=65537, key_size=2048)


def ecdsa_private_key():
    """Generate a new ECDSA private key using SECP256R1 curve."""
    return ec.generate_private_key(ec.SECP256R1(), default_backend())


def pem_public_key(private_key: PrivateKeyTypes) -> bytes:
    """Return the public key in PEM format from a private key."""
    return private_key.public_key().public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo,
    )


def save_private_key(cid: str, private_key: PrivateKeyTypes):
    """Save the private key to a file in the rembus directory."""
    kdir = os.path.join(rs.rembus_dir(), cid)

    if not os.path.exists(kdir):
        os.makedirs(kdir)

    fn = os.path.join(kdir, ".secret")
    private_key_file = open(fn, "wb")

    pem_private_key = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.TraditionalOpenSSL,
        encryption_algorithm=serialization.NoEncryption(),
    )

    private_key_file.write(pem_private_key)
    private_key_file.close()


def load_private_key(cid: str) -> PrivateKeyTypes:
    """Load the private key from a file in the rembus directory."""
    fn = os.path.join(rs.rembus_dir(), cid, ".secret")
    # fn = os.path.join(rembus_dir(), cid)
    with open(fn, "rb") as key_file:
        private_key = serialization.load_pem_private_key(
            key_file.read(), password=None
        )

    return private_key


def load_public_key(router, cid: str):
    """Load the public key from a file in the rembus directory."""
    fn = rs.key_file(router.id, cid)
    try:
        with open(fn, "rb") as f:
            public_key = serialization.load_pem_public_key(
                f.read(),
            )
    except ValueError:
        try:
            with open(fn, "rb") as f:
                public_key = serialization.load_der_public_key(
                    f.read(),
                )
        except ValueError as e:
            raise ValueError(
                f"Could not load public key from file: {fn}"
            ) from e

    return public_key


def save_pubkey(router_id: str, cid: str, pubkey: bytes, keytype: int):
    """Save the public key to a file in the rembus directory."""
    name = rs.key_base(router_id, cid)
    keyformat = "der"
    # check if pubkey start with -----BEGIN chars
    if pubkey[0:10] == bytes(
        [0x2D, 0x2D, 0x2D, 0x2D, 0x2D, 0x42, 0x45, 0x47, 0x49, 0x4E]
    ):
        keyformat = "pem"

    if keytype == SIG_RSA:
        typestr = "rsa"
    else:
        typestr = "ecdsa"

    fn = f"{name}.{typestr}.{keyformat}"
    with open(fn, "wb") as f:
        f.write(pubkey)


def remove_pubkey(router, cid: str):
    """Remove the public key file for a component."""
    fn = rs.key_file(router.id, cid)
    os.remove(fn)
