import zmq
import random
import socket as sock
from typing import Optional, Dict, Any, Callable
from arcros.database import DatabaseManager


def _get_subscriber_id() -> str:
    hostname = sock.gethostname()
    pid = __import__('os').getpid()
    return f"{hostname}_{pid}"


class Publisher:
    def __init__(self, name: str, port: Optional[int] = None, host: str = 'localhost', description: str = ''):
        self.name = name
        self.host = host
        self.db = DatabaseManager()
        self.context = zmq.Context()
        self.socket = self.context.socket(zmq.PUB)
        
        existing = self.db.get_publisher(name)
        if existing:
            self.port = existing['port']
            self.host = existing['host']
        else:
            if port is None:
                self.port = self._find_available_port()
            else:
                self.port = port
            
            self.db.add_publisher(name, self.port, self.host, description)
        
        self.socket.bind(f"tcp://{self.host}:{self.port}")
        print(f"Publisher '{name}' bound to {self.host}:{self.port}")
    
    def _find_available_port(self) -> int:
        port = self.db.find_available_port(self.host)
        if port is None:
            raise RuntimeError("No available ports in configured range. Use 'zmq-monitor config' to adjust port range.")
        
        try:
            test_socket = self.context.socket(zmq.PUB)
            test_socket.bind(f"tcp://{self.host}:{port}")
            test_socket.close()
            return port
        except zmq.ZMQError:
            raise RuntimeError(f"Port {port} is in use by another process")
    
    def publish(self, data: Dict[str, Any]):
        self.socket.send_json(data)
    
    def close(self):
        self.socket.close()
        self.context.term()
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


class Subscriber:
    def __init__(self, name_or_port: str, host: str = 'localhost', subscriber_name: Optional[str] = None):
        self.db = DatabaseManager()
        self.context = zmq.Context()
        self.socket = self.context.socket(zmq.SUB)
        self.socket.setsockopt(zmq.SUBSCRIBE, b'')
        
        self.subscriber_name = subscriber_name or _get_subscriber_id()
        self.publisher_name = None
        
        if name_or_port.isdigit():
            port = int(name_or_port)
            self.publisher_name = f"port_{port}"
        else:
            info = self.db.get_publisher(name_or_port)
            if not info:
                raise ValueError(f"Publisher '{name_or_port}' not found in database")
            port = info['port']
            host = info['host']
            self.publisher_name = name_or_port
        
        self.host = host
        self.port = port
        self.socket.connect(f"tcp://{host}:{port}")
        
        if self.publisher_name and not self.publisher_name.startswith("port_"):
            self.db.add_subscription(self.subscriber_name, self.publisher_name, sock.gethostname())
        
        print(f"Subscriber '{self.subscriber_name}' connected to '{self.publisher_name}' at {host}:{port}")
    
    def receive(self, timeout_ms: Optional[int] = None) -> Optional[Dict[str, Any]]:
        if timeout_ms:
            self.socket.setsockopt(zmq.RCVTIMEO, timeout_ms)
        else:
            self.socket.setsockopt(zmq.RCVTIMEO, -1)
        
        try:
            return self.socket.recv_json()
        except zmq.Again:
            return None
    
    def receive_nonblocking(self) -> Optional[Dict[str, Any]]:
        try:
            return self.socket.recv_json(flags=zmq.NOBLOCK)
        except zmq.Again:
            return None
    
    def listen(self, callback: Callable[[Dict[str, Any]], None]):
        print(f"Listening to '{self.publisher_name}' at {self.host}:{self.port}")
        print("Press Ctrl+C to stop\n")
        
        try:
            while True:
                try:
                    data = self.socket.recv_json(flags=zmq.NOBLOCK)
                    callback(data)
                except zmq.Again:
                    import time
                    time.sleep(0.01)
        except KeyboardInterrupt:
            print("\nStopped listening.")
    
    def close(self):
        if self.publisher_name and not self.publisher_name.startswith("port_"):
            self.db.remove_subscription(self.subscriber_name, self.publisher_name)
        self.socket.close()
        self.context.term()
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


def create_publisher(name: str, port: Optional[int] = None, host: str = 'localhost', description: str = '') -> Publisher:
    return Publisher(name, port, host, description)


def create_subscriber(name_or_port: str, host: str = 'localhost', subscriber_name: Optional[str] = None) -> Subscriber:
    return Subscriber(name_or_port, host, subscriber_name)