from feathersdk.comms.comms_manager import CommsManager, CommsManagerNotRunningError, SocketResult, SendCanMessage, \
    SendTCPMessage, SendCallbackParamType
from feathersdk.comms.system import get_can_links_json, is_can_interface
from feathersdk.utils.feathertypes import Union, Optional, get_union_types
from dataclasses import dataclass, field
from feathersdk.utils.common import FeatherThreadLock, currtime, timediff, timestamped_class, TimeType, listdirfull
from feathersdk.utils.telemetry import CircularBuffer
import os
import gzip
import threading
import time
import socket


WIRELOG_FILE_EXTENSION = ".wirelog"


@dataclass
class WireHawkConfig:
    """Configuration for WireHawk."""
    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"""
    can_interface_refresh_interval_seconds: float = field(default=0.3)
    """How often to check for new/removed CAN/vCAN interfaces."""


@dataclass
@timestamped_class
class WireHawkSnapshot:
    """A snapshot of the statistics for a single endpoint."""
    endpoint: str
    """The endpoint name."""
    endpoint_uid: int
    """The unique identifier for the 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."""
    endpoint_uid: int
    """The unique identifier for the endpoint."""
    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, ep: str, ep_uid: int, config: WireHawkConfig):
        self.config = config
        self.endpoint = ep
        self.endpoint_uid = ep_uid

        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(
            endpoint=self.endpoint,
            endpoint_uid=self.endpoint_uid,
            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: Union[SendCallbackParamType, SocketResult]) -> None:
        """Add a message to the buffers"""
        if isinstance(message, get_union_types(SendCallbackParamType)):
            self.sent_messages.append(message)
        elif isinstance(message, SocketResult):
            self.recv_messages.append(message)
        else:
            raise NotImplementedError(f"Unknown message type: {type(message)}")


def _ep_uid(ep: str) -> int:
    """Get the unique identifier for an endpoint, resilient to deleting/recreating the interface with the same name.
    
    Assumes the endpoint does in fact exist.
    """
    return socket.if_nametoindex(ep)


def _curr_eps() -> list[str]:
    """Get the current open CAN/vCAN interface names"""
    return {link.get("ifname", "") for link in get_can_links_json()}


# Global dictionary of all WireHawk instances by their unique identifier. Really only used for testing.
# NOTE: we add ourselves to it automatically, but this must be cleaned up manually by the caller.
_WIREHAWK_INSTANCES: dict[int, "WireHawk"] = {}


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)

        eps = _curr_eps()
        self.comms.add_endpoint(eps, skip_duplicates=True)

        self._lock = FeatherThreadLock("WireHawk")

        self.stats = {_ep_uid(ep): WireHawkEndpoint(ep=ep, ep_uid=_ep_uid(ep), config=self.config) for ep in eps}
        self.wirehawk_init_time = currtime()
        self.last_log_time = currtime()
        self.log_file_idx = 0
        self._known_can_ep_uids = set(_ep_uid(ep) for ep in eps if is_can_interface(ep))

        if self.config.log_dir is not None and not os.path.exists(self.config.log_dir):
            os.makedirs(self.config.log_dir)
        
        self._running_refresh = True
        self._refresh_thread = threading.Thread(target=self.__refresh_can_endpoints_thread, daemon=True)
        self._refresh_thread.start()

        _WIREHAWK_INSTANCES[id(self)] = self

    def on_recv(self, message: SocketResult):
        with self._lock:
            self._ensure_endpoint_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._ensure_endpoint_stats(message.iface).add_message(message)
            elif isinstance(message, SendTCPMessage):
                self._ensure_endpoint_stats(message.ip).add_message(message)
            else:
                raise NotImplementedError(f"Unknown message type: {type(message)}")
            self._possibly_flush_logs()
    
    def _ensure_endpoint_stats(self, ep: str) -> WireHawkEndpoint:
        """Make sure we have a stats entry for the endpoint. Assumes we have the lock already."""
        assert self._lock.is_locked()
        try:
            ep_uid = _ep_uid(ep)
        except OSError:
            return WireHawkEndpoint(ep=ep, ep_uid=0, config=self.config)  # In case the interface was deleted
        if ep_uid not in self.stats:
            self.stats[ep_uid] = WireHawkEndpoint(ep=ep, ep_uid=ep_uid, config=self.config)
        return self.stats[ep_uid]
    
    def __refresh_can_endpoints_thread(self) -> None:
        """Detect new/removed CAN/vCAN interfaces and update comms subscriptions.
        
        IMPORTANT:
        - We never remove entries from self.stats, even if an interface disappears.
        - We only restart CommsManager when there is an actual interface set change.
        """
        last_refresh_time = 0
        while True:
            # Check if we are ready to refresh yet
            if timediff(last_refresh_time) < self.config.can_interface_refresh_interval_seconds:
                time.sleep(min(self.config.can_interface_refresh_interval_seconds - timediff(last_refresh_time), 0.01))
                continue
                
            # Update the last refresh time and check we should still be running
            last_refresh_time = currtime()
            if not self._running_refresh:
                break
                
            # Check that the CommsManager is still running. If not, assume user closed it and we should stop.
            if not self.comms.is_running():
                break

            # Get the current list of CAN/vCAN interfaces, both in OS and in CommsManager
            curr_eps: dict[str, int] = {}
            for name in set(link.get("ifname", "") for link in get_can_links_json()):
                try:
                    curr_eps[name] = _ep_uid(name)
                except OSError:
                    pass
            
            curr_comms_ep_uids: set[int] = set()
            for ep in self.comms.endpoints:
                try:
                    curr_comms_ep_uids.add(_ep_uid(ep))
                except OSError:
                    pass

            # Update if either there are new interfaces, or the CommsManager was cut down to a smaller number of them
            if set(curr_eps.values()) != self._known_can_ep_uids or curr_comms_ep_uids != self._known_can_ep_uids:
                with self._lock:
                    # Add any new ep_uids if they are now running
                    for ep in curr_eps:
                        self._ensure_endpoint_stats(ep)
                        if ep in self.comms.endpoints:
                            self.comms.remove_endpoint(ep)  # Make sure we only remove changed endpoints
                    
                        try:  # Sometimes errors at the end of tests, and I don't want the error message to print
                            self.comms.add_endpoint(list(curr_eps), enable_cans=False, allow_no_enable_can=True)
                        except OSError:
                            del self.stats[curr_eps[ep]]
                            continue
                    
                    # Remove any old ep_uids if they are no longer running, and update the known ep_uids
                    del_uids = set(self.stats.keys()) - set(curr_eps.values())
                    for ep_uid in del_uids:
                        del self.stats[ep_uid]
                    self._known_can_ep_uids = set(curr_eps.values())
            
            time.sleep(self.config.can_interface_refresh_interval_seconds)
    
    def __clear_stats(self) -> None:
        """Clear the stats. Assumes we have the lock already."""
        assert self._lock.is_locked()
        for ep_stat in self.stats.values():
            ep_stat.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:
            # Still keep in-memory stats for live inspection. CircularBuffer capacity bounds memory.
            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 and \
            not any(ep_uid not in self._known_can_ep_uids for ep_uid in self.stats.keys()):
            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_stat in self.stats.values():
                buf: Union[CircularBuffer[SendCallbackParamType], CircularBuffer[SocketResult]]  # type hint for loop
                for buf in [ep_stat.sent_messages, ep_stat.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))
        
        # Remove any old ep_uids if they are no longer running
        for ep_uid in self.stats.keys():
            if ep_uid not in self._known_can_ep_uids:
                del self.stats[ep_uid]
    
    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)
    
    def close(self) -> None:
        """Close the WireHawk instance. Freezes the stats and stops the refresh thread. Also closes the CommsManager."""
        with self._lock:
            self._running_refresh = False
            if self._refresh_thread.is_alive():
                self._refresh_thread.join(timeout=5.0)
                if self._refresh_thread.is_alive():
                    raise RuntimeError("Refresh thread did not join in time")
            
            try:
                self.comms.reset()
            except CommsManagerNotRunningError:
                pass
    