from __future__ import annotations

import enum
import queue
import threading
import time
import typing
import uuid
from types import TracebackType
from typing import Type, Callable, TypeVar, Tuple
import RNS
from RNS.vendor import umsgpack
from RNS.Buffer import StreamDataMessage as RNSStreamDataMessage
import rnsh.retry
import abc
import contextlib
import struct
import logging as __logging
from abc import ABC, abstractmethod

module_logger = __logging.getLogger(__name__)

MSG_MAGIC = 0xac
PROTOCOL_VERSION = 1


def _make_MSGTYPE(val: int):
    return ((MSG_MAGIC << 8) & 0xff00) | (val & 0x00ff)


class NoopMessage(RNS.MessageBase):
    MSGTYPE = _make_MSGTYPE(0)

    def pack(self) -> bytes:
        return bytes()

    def unpack(self, raw):
        pass


class WindowSizeMessage(RNS.MessageBase):
    MSGTYPE = _make_MSGTYPE(2)

    def __init__(self, rows: int = None, cols: int = None, hpix: int = None, vpix: int = None):
        super().__init__()
        self.rows = rows
        self.cols = cols
        self.hpix = hpix
        self.vpix = vpix

    def pack(self) -> bytes:
        return umsgpack.packb((self.rows, self.cols, self.hpix, self.vpix))

    def unpack(self, raw):
        self.rows, self.cols, self.hpix, self.vpix = umsgpack.unpackb(raw)


class ExecuteCommandMesssage(RNS.MessageBase):
    MSGTYPE = _make_MSGTYPE(3)

    def __init__(self, cmdline: [str] = None, pipe_stdin: bool = False, pipe_stdout: bool = False,
                 pipe_stderr: bool = False, tcflags: [any] = None, term: str | None = None, rows: int = None,
                 cols: int = None, hpix: int = None, vpix: int = None):
        super().__init__()
        self.cmdline = cmdline
        self.pipe_stdin = pipe_stdin
        self.pipe_stdout = pipe_stdout
        self.pipe_stderr = pipe_stderr
        self.tcflags = tcflags
        self.term = term
        self.rows = rows
        self.cols = cols
        self.hpix = hpix
        self.vpix = vpix

    def pack(self) -> bytes:
        return umsgpack.packb((self.cmdline, self.pipe_stdin, self.pipe_stdout, self.pipe_stderr,
                               self.tcflags, self.term, self.rows, self.cols, self.hpix, self.vpix))

    def unpack(self, raw):
        self.cmdline, self.pipe_stdin, self.pipe_stdout, self.pipe_stderr, self.tcflags, self.term, self.rows, \
            self.cols, self.hpix, self.vpix = umsgpack.unpackb(raw)


# Create a version of RNS.Buffer.StreamDataMessage that we control
class StreamDataMessage(RNSStreamDataMessage):
    MSGTYPE = _make_MSGTYPE(4)
    STREAM_ID_STDIN  = 0
    STREAM_ID_STDOUT = 1
    STREAM_ID_STDERR = 2


class VersionInfoMessage(RNS.MessageBase):
    MSGTYPE = _make_MSGTYPE(5)

    def __init__(self, sw_version: str = None):
        super().__init__()
        self.sw_version = sw_version or rnsh.__version__
        self.protocol_version = PROTOCOL_VERSION

    def pack(self) -> bytes:
        return umsgpack.packb((self.sw_version, self.protocol_version))

    def unpack(self, raw):
        self.sw_version, self.protocol_version = umsgpack.unpackb(raw)


class ErrorMessage(RNS.MessageBase):
    MSGTYPE = _make_MSGTYPE(6)

    def __init__(self, msg: str = None, fatal: bool = False, data: dict = None):
        super().__init__()
        self.msg = msg
        self.fatal = fatal
        self.data = data

    def pack(self) -> bytes:
        return umsgpack.packb((self.msg, self.fatal, self.data))

    def unpack(self, raw: bytes):
        self.msg, self.fatal, self.data = umsgpack.unpackb(raw)


class CommandExitedMessage(RNS.MessageBase):
    MSGTYPE = _make_MSGTYPE(7)

    def __init__(self, return_code: int = None):
        super().__init__()
        self.return_code = return_code

    def pack(self) -> bytes:
        return umsgpack.packb(self.return_code)

    def unpack(self, raw: bytes):
        self.return_code = umsgpack.unpackb(raw)


message_types = [NoopMessage, VersionInfoMessage, WindowSizeMessage, ExecuteCommandMesssage, StreamDataMessage,
                 CommandExitedMessage, ErrorMessage]


def register_message_types(channel: RNS.Channel.Channel):
    for message_type in message_types:
        channel.register_message_type(message_type)