# This file is part of the Trezor project.
#
# Copyright (C) SatoshiLabs and contributors
#
# This library is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License version 3
# as published by the Free Software Foundation.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the License along with this library.
# If not, see <https://www.gnu.org/licenses/lgpl-3.0.html>.

from __future__ import annotations

import struct

from typing_extensions import Self

CODEC_V1 = 0x3F
CONTINUATION_PACKET = 0x80
HANDSHAKE_INIT_REQ = 0x00
HANDSHAKE_INIT_RES = 0x01
HANDSHAKE_COMP_REQ = 0x02
HANDSHAKE_COMP_RES = 0x03
ENCRYPTED_TRANSPORT = 0x04

CONTINUATION_PACKET_MASK = 0x80
ACK_MASK = 0xF7
DATA_MASK = 0xE7

ACK_MESSAGE = 0x20
_ERROR = 0x42
CHANNEL_ALLOCATION_REQ = 0x40
_CHANNEL_ALLOCATION_RES = 0x41

PING = 0x43
PONG = 0x44

TREZOR_STATE_UNPAIRED = b"\x00"
TREZOR_STATE_PAIRED = b"\x01"

BROADCAST_CHANNEL_ID = 0xFFFF


class MessageHeader:
    format_str_init = ">BHH"
    format_str_cont = ">BH"

    def __init__(self, ctrl_byte: int, cid: int, length: int) -> None:
        self.ctrl_byte = ctrl_byte
        self.cid = cid
        self.data_length = length

    def to_bytes_init(self) -> bytes:
        return struct.pack(
            self.format_str_init, self.ctrl_byte, self.cid, self.data_length
        )

    def to_bytes_cont(self) -> bytes:
        return struct.pack(self.format_str_cont, CONTINUATION_PACKET, self.cid)

    def pack_to_init_buffer(self, buffer: bytearray, buffer_offset: int = 0) -> None:
        struct.pack_into(
            self.format_str_init,
            buffer,
            buffer_offset,
            self.ctrl_byte,
            self.cid,
            self.data_length,
        )

    def pack_to_cont_buffer(self, buffer: bytearray, buffer_offset: int = 0) -> None:
        struct.pack_into(
            self.format_str_cont, buffer, buffer_offset, CONTINUATION_PACKET, self.cid
        )

    def is_ack(self) -> bool:
        return self.ctrl_byte & ACK_MASK == ACK_MESSAGE

    def is_channel_allocation_response(self) -> bool:
        return (
            self.cid == BROADCAST_CHANNEL_ID
            and self.ctrl_byte == _CHANNEL_ALLOCATION_RES
        )

    def is_pong(self) -> bool:
        return self.cid == BROADCAST_CHANNEL_ID and self.ctrl_byte == PONG

    def is_handshake_init_response(self) -> bool:
        return self.ctrl_byte & DATA_MASK == HANDSHAKE_INIT_RES

    def is_handshake_comp_response(self) -> bool:
        return self.ctrl_byte & DATA_MASK == HANDSHAKE_COMP_RES

    def is_encrypted_transport(self) -> bool:
        return self.ctrl_byte & DATA_MASK == ENCRYPTED_TRANSPORT

    @classmethod
    def get_error_header(cls, cid: int, length: int) -> Self:
        return cls(_ERROR, cid, length)

    @classmethod
    def get_channel_allocation_request_header(cls, length: int) -> Self:
        return cls(CHANNEL_ALLOCATION_REQ, BROADCAST_CHANNEL_ID, length)

    @classmethod
    def get_ping_header(cls, length: int) -> Self:
        return cls(PING, BROADCAST_CHANNEL_ID, length)
