import socket
import threading
from typing import Callable

class PeerConnection:
    """
    Represents a live, two-way connection to a peer.
    We keep a socket open, read data in a background thread,
    and can send data from the main thread at any time.
    """
    def __init__(self, sock: socket.socket, on_data_received: Callable[['PeerConnection', bytes], None]):
        self.sock = sock
        self.on_data_received = on_data_received
        self.alive = True

        self.read_thread = threading.Thread(target=self.read_loop, daemon=True)
        self.read_thread.start()

    def read_loop(self):
        while self.alive:
            try:
                data = self.sock.recv(4096)
                if not data:
                    break
                self.on_data_received(self, data)
            except OSError:
                break
        self.alive = False
        self.sock.close()
        print("[PeerConnection] Connection closed.")

    def send(self, data: bytes):
        if not self.alive:
            print("[PeerConnection.send] Cannot send, connection not alive.")
            return
        try:
            self.sock.sendall(data)
        except OSError:
            print("[PeerConnection.send] Send failed, connection might be closed.")
            self.alive = False

    def close(self):
        self.alive = False
        try:
            self.sock.shutdown(socket.SHUT_RDWR)
        except OSError:
            pass
        self.sock.close()


class ServerListener(threading.Thread):
    """
    A thread that listens on a given port. When a new client connects,
    it creates a PeerConnection for that client.
    """
    def __init__(self, host: str, port: int,
                 on_new_connection: Callable[[PeerConnection], None],
                 on_data_received: Callable[[PeerConnection, bytes], None]):
        super().__init__(daemon=True)
        self.host = host
        self.port = port
        self.on_new_connection = on_new_connection
        self.on_data_received = on_data_received
        self.server_socket = None
        self.stop_event = threading.Event()

    def run(self):
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.bind((self.host, self.port))
        self.server_socket.listen(5)
        self.server_socket.settimeout(1.0)
        print(f"[ServerListener] Listening on {self.host}:{self.port}")

        while not self.stop_event.is_set():
            try:
                client_sock, addr = self.server_socket.accept()
                print(f"[ServerListener] Accepted connection from {addr}")
                conn = PeerConnection(client_sock, self.on_data_received)
                self.on_new_connection(conn)
            except socket.timeout:
                pass
            except OSError:
                break

        if self.server_socket:
            self.server_socket.close()

    def stop(self):
        self.stop_event.set()
        if self.server_socket:
            self.server_socket.close()


def connect_to_peer(host: str, port: int,
                    on_data_received: Callable[[PeerConnection, bytes], None]) -> PeerConnection:
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.connect((host, port))
    print(f"[connect_to_peer] Connected to {host}:{port}")
    conn = PeerConnection(sock, on_data_received)
    return conn
