from feathersdk.comms.comms_manager import CommsManager, SocketResult, SendCanMessage, SendTCPMessage, \
    SendCallbackParamType
from feathersdk.utils.feathertypes import Union, Optional
from dataclasses import dataclass, field
from feathersdk.utils.common import FeatherThreadLock, currtime, timediff, TimestampedMixin, TimeType, listdirfull
from feathersdk.utils.telemetry import CircularBuffer
import os
import gzip


WIRELOG_FILE_EXTENSION = ".wirelog"


@dataclass
class WireHawkConfig:
    """Configuration for WireHawk."""
    endpoints: list[str]
    """The endpoints to log messages from."""
    log_dir: Optional[str] = field(default_factory=lambda: f"/feathersys/logs/wirehawk/{currtime()}")
    """The directory to log messages to. Will create one subdirectory per endpoint. Set None to disable file logging."""
    max_file_size: int = field(default=50 * 2**20)  # 50 MB
    """The maximum size of a single log file before compressing it and starting a new one."""
    max_dir_size: int = field(default=2**30)  # 1 GB
    """The maximum size of an endpoint's subdirectory before deleting the oldest compressed file."""
    log_interval_seconds: float = field(default=2.0)
    """Maximum amount of time to wait between flushing the message buffers to file"""
    max_message_buffer_size: int = field(default=1000)
    """Maximum number of messages to buffer before flushing to file"""
    send_stats_recent_messages_size: int = field(default=1000)
    """Max number of most recent messages to return when getting stats snapshot"""
    recv_stats_recent_messages_size: int = field(default=1000)
    """Max number of most recent messages to return when getting stats snapshot"""


@dataclass
class WireHawkSnapshot(TimestampedMixin):
    """A snapshot of the statistics for a single endpoint."""
    avg_send_rate: float
    """The average messages-sent per second of the endpoint."""
    avg_recv_rate: float
    """The average messages-recv'd per second of the endpoint."""
    latest_send_time: TimeType
    """The time of the last message that was sent."""
    latest_recv_time: TimeType
    """The time of the last message that was received."""
    earliest_send_time: TimeType
    """The time of the earliest message in our buffer that was sent."""
    earliest_recv_time: TimeType
    """The time of the earliest message in our buffer that was received."""
    most_recent_sent_messages: Union[list[SendCanMessage], list[SendTCPMessage]]
    """The most recent messages that were sent."""
    most_recent_recv_messages: list[SocketResult]
    """The most recent messages that were received."""


class WireHawkEndpoint:
    """Statistics for a single endpoint."""
    endpoint: str
    """The endpoint name."""
    sent_messages: Union[list[SendCanMessage], list[SendTCPMessage]]
    """The messages that have been sent."""
    recv_messages: list[SocketResult]
    """The messages that have been received."""

    def __init__(self, endpoint: str, config: WireHawkConfig):
        self.config = config
        self.endpoint = endpoint

        ms = self.config.max_message_buffer_size
        self.sent_messages: CircularBuffer[SendCanMessage | SendTCPMessage] = CircularBuffer(capacity=ms, dtype=object)
        self.recv_messages: CircularBuffer[SocketResult] = CircularBuffer(capacity=ms, dtype=object)
    
    def get_last_stats_snapshot(self) -> WireHawkSnapshot:
        """Get the last stats snapshot for this endpoint."""
        return WireHawkSnapshot(
            avg_send_rate=sum(msg.timestamp for msg in self.sent_messages) / max(len(self.sent_messages), 1),
            avg_recv_rate=sum(msg.timestamp for msg in self.recv_messages) / max(len(self.recv_messages), 1),
            latest_send_time=self.sent_messages[-1].timestamp if self.sent_messages.count > 0 else -1,
            latest_recv_time=self.recv_messages[-1].timestamp if self.recv_messages.count > 0 else -1,
            earliest_send_time=self.sent_messages[0].timestamp if self.sent_messages.count > 0 else -1,
            earliest_recv_time=self.recv_messages[0].timestamp if self.recv_messages.count > 0 else -1,
            most_recent_sent_messages=self.sent_messages[-self.config.send_stats_recent_messages_size:],
            most_recent_recv_messages=self.recv_messages[-self.config.recv_stats_recent_messages_size:],
        )
    
    def clear_buffers(self) -> None:
        """Clear the buffers"""
        self.sent_messages.clear()
        self.recv_messages.clear()
    
    def add_message(self, message: SendCallbackParamType | SocketResult) -> None:
        """Add a message to the buffers"""
        if isinstance(message, SendCallbackParamType):
            self.sent_messages.append(message)
        elif isinstance(message, SocketResult):
            self.recv_messages.append(message)
        else:
            raise NotImplementedError(f"Unknown message type: {type(message)}")


class WireHawk:
    """Like WireShark, but also nothing like WireShark.

    Feather tool to intercept and log comms messages. Mostly for debugging. Will begin listening on initialization.
    """

    def __init__(self, config: WireHawkConfig):
        self.config = config
        self.comms = CommsManager()
        self.comms.add_callback(self.on_recv)
        self.comms._add_send_callback(self.on_send)
        self.comms.restart_with_added_endpoints(self.config.endpoints, start_ok=True)

        self._lock = FeatherThreadLock("WireHawk")

        self.stats = {ep: WireHawkEndpoint(endpoint=ep, config=self.config) for ep in self.config.endpoints}
        self.wirehawk_init_time = currtime()
        self.last_log_time = currtime()
        self.log_file_idx = 0

        if not os.path.exists(self.config.log_dir):
            os.makedirs(self.config.log_dir)

    def on_recv(self, message: SocketResult):
        with self._lock:
            self.stats[message.socket_name].add_message(message)
            self._possibly_flush_logs()

    def on_send(self, message: SendCallbackParamType):
        with self._lock:
            if isinstance(message, SendCanMessage):
                self.stats[message.iface].add_message(message)
            elif isinstance(message, SendTCPMessage):
                self.stats[message.ip].add_message(message)
            else:
                raise NotImplementedError(f"Unknown message type: {type(message)}")
            self._possibly_flush_logs()
    
    def __clear_stats(self) -> None:
        """Clear the stats. Assumes we have the lock already."""
        assert self._lock.is_locked()
        for ep in self.stats.values():
            ep.clear_buffers()
    
    def _possibly_flush_logs(self) -> None:
        """Check if we should flush logs to files or not and log if so. Assumes we have the lock already."""
        assert self._lock.is_locked()
        if self.config.log_dir is None:
            self.__clear_stats()
            return
        
        # Return now if we don't have enough time passed or messages buffered since the last log
        num_messages = sum(len(ep.sent_messages) + len(ep.recv_messages) for ep in self.stats.values())
        if timediff(self.last_log_time) < self.config.log_interval_seconds and \
            num_messages < self.config.max_message_buffer_size:
            return
        
        self.__flush_logs()
        
        self.__clear_stats()
        self.last_log_time = currtime()
    
    def __flush_logs(self) -> None:
        """Flush logs to files. Assumes we have the lock already and we are logging to files."""
        assert self._lock.is_locked()
        
        # Build the list of messages to log in the order they were sent/received
        messages_to_log: list[tuple[TimeType, Union[SendCallbackParamType, SocketResult]]] = []
        while True:
            oldest_time: Optional[int] = None
            oldest_buf: Optional[Union[CircularBuffer[SendCallbackParamType], CircularBuffer[SocketResult]]] = None

            for ep in self.stats.values():
                buf: Union[CircularBuffer[SendCallbackParamType], CircularBuffer[SocketResult]]  # type hint for loop
                for buf in [ep.sent_messages, ep.recv_messages]:
                    if len(buf) > 0 and (oldest_time is None or buf[0].timestamp < oldest_time):
                        oldest_time = buf[0].timestamp
                        oldest_buf = buf

            if oldest_time is None:
                break
            
            messages_to_log.append((oldest_time, oldest_buf.pop(0)))
        
        # Append the messages to the current log file
        file_path = os.path.join(self.config.log_dir, f"wirehawk_{self.log_file_idx}{WIRELOG_FILE_EXTENSION}")
        with open(file_path, "a") as f:
            for time, message in messages_to_log:
                f.write(f"[{(self.wirehawk_init_time - time):.6f}s] - {repr(message)}\n")
        
        # Change the permissions of the log file to 777 in case we are running as a non-root user
        os.chmod(file_path, 0o777)
        
        # Check if the log file is too large, and if so, compress it and start a new one
        if os.path.getsize(file_path) > self.config.max_file_size:
            self.__compress_log_file(file_path)
            self.log_file_idx += 1
        
        # Check if the log directory is too large, and if so, delete the oldest compressed file
        wirelog_files = listdirfull(self.config.log_dir, endswith=(WIRELOG_FILE_EXTENSION, ".gz"))
        if sum(os.path.getsize(f) for f in wirelog_files) > self.config.max_dir_size:
            os.remove(min(wirelog_files, key=os.path.getctime))
    
    def __compress_log_file(self, file_path: str) -> None:
        """Compress the log file. Assumes we have the lock already and we are logging to files."""
        assert self._lock.is_locked()
        
        new_file_path = file_path + ".gz"
        with open(file_path, "rb") as f:
            data = f.read()
        with gzip.open(new_file_path, "wb") as f:
            f.write(data)
        
        os.chmod(new_file_path, 0o777)
        os.remove(file_path)