import socket
import struct
import threading
from typing import Callable, Optional

import wx

from .protocol import NetworkMesssage

OnMessageCallback = Callable[[NetworkMesssage], None]
OnNoticeCallback = Callable[[str], None]


def receive_exact(sock: socket.socket, n: int) -> Optional[bytes]:
    data = bytearray()

    while len(data) < n:
        packet = sock.recv(n - len(data))
        if not packet:
            return None
        data.extend(packet)

    return bytes(data)


class NetworkController(threading.Thread):
    def __init__(
        self,
        on_message_callback: OnMessageCallback,
        on_notice_callback: Optional[OnNoticeCallback] = None,
    ):
        super().__init__(daemon=True)
        self.on_message_callback = on_message_callback
        self.on_notice_callback = on_notice_callback

        self.sock: Optional[socket.socket] = None
        self.conn: Optional[socket.socket] = None

        self.is_running: bool = False
        self.is_server: bool = False

    def serve(self, host: str = "0.0.0.0", port: int = 8994) -> None:
        self.is_server = True

        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.sock.bind((host, port))
        self.sock.listen(1)

        self.is_running = True
        self.start()

    def connect(self, host: str, port: int = 8994) -> bool:
        self.is_server = False

        try:
            self.conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.conn.connect((host, port))

            self.is_running = True
            self._notify("Connected to host")
            self.start()
            return True
        except Exception as e:
            self._notify(f"Connection failed: {e}")
            return False

    def send(self, event: str, payload: dict = {}) -> None:
        if not self.conn or not self.is_running:
            return

        m = NetworkMesssage(event, payload)

        try:
            self.conn.sendall(m.encode())
        except Exception as e:
            self._notify(f"Send error: {e}")
            self._close_active_connection()

    def run(self) -> None:
        HEADER_SIZE = 4

        while self.is_running:
            if self.is_server and self.sock and not self.conn:
                self._notify("Waiting for peer to connect...")

                try:
                    self.conn, addr = self.sock.accept()
                    self._notify(f"Peer connected from {addr[0]}:{addr[1]}")
                except Exception:
                    break

            while self.is_running and self.conn:
                try:
                    header_bytes = receive_exact(self.conn, HEADER_SIZE)
                    if not header_bytes:
                        self._notify("Peer disconnected")
                        break

                    body_size = struct.unpack(">I", header_bytes)[0]

                    body_bytes = receive_exact(self.conn, body_size)
                    if not body_bytes:
                        self._notify("Peer disconnected unexpectedly")
                        break

                    json_str = body_bytes.decode("utf-8")
                    m = NetworkMesssage.from_json(json_str)

                    wx.CallAfter(self.on_message_callback, m)

                except Exception as e:
                    self._notify(f"Read error: {e}")
                    break

            self._close_active_connection()

            if not self.is_server:
                break

        self.stop()

    def _close_active_connection(self) -> None:
        if self.conn:
            try:
                self.conn.shutdown(socket.SHUT_RDWR)
                self.conn.close()
            except Exception:
                pass
            finally:
                self.conn = None

    def _notify(self, notice: str) -> None:
        if self.on_notice_callback:
            wx.CallAfter(self.on_notice_callback, notice)

    def stop(self) -> None:
        self.is_running = False

        self._close_active_connection()

        if self.sock:
            try:
                self.sock.close()
            except Exception:
                pass
            finally:
                self.sock = None

        self._notify("Connection closed")