import zmq
import argparse
import time
from datetime import datetime
from arcros.database import DatabaseManager


class ZMQMonitor:
    def __init__(self):
        self.context = zmq.Context()
        self.db = DatabaseManager()
    
    def discover_and_list(self, host='localhost', timeout=500):
        registered_publishers = self.db.list_all_publishers()
        
        if not registered_publishers:
            print("No publishers registered in database.")
            print("Use 'create_publisher()' to create publishers.")
            return
        
        print(f"Checking {len(registered_publishers)} registered publisher(s)...\n")
        
        active_pubs = []
        
        for pub_info in registered_publishers:
            name = pub_info['name']
            port = pub_info['port']
            pub_host = pub_info['host']
            
            if pub_host != host:
                continue
            
            socket = self.context.socket(zmq.SUB)
            socket.setsockopt(zmq.SUBSCRIBE, b'')
            socket.setsockopt(zmq.RCVTIMEO, timeout)
            
            try:
                socket.connect(f"tcp://{pub_host}:{port}")
                data = socket.recv_json()
                
                if isinstance(data, dict):
                    data_types = {k: type(v).__name__ for k, v in data.items()}
                else:
                    data_types = type(data).__name__
                
                active_pubs.append({
                    'name': name,
                    'port': port,
                    'host': pub_host,
                    'data_types': data_types,
                    'status': 'active'
                })
                
            except zmq.Again:
                active_pubs.append({
                    'name': name,
                    'port': port,
                    'host': pub_host,
                    'data_types': '-',
                    'status': 'inactive'
                })
            except Exception:
                active_pubs.append({
                    'name': name,
                    'port': port,
                    'host': pub_host,
                    'data_types': '-',
                    'status': 'error'
                })
            finally:
                socket.close()
        
        if not active_pubs:
            print(f"No publishers found on host '{host}'.")
            return
        
        print(f"{'Name':<20} {'Port':<8} {'Host':<15} {'Status':<10} {'Data Types'}")
        print("-" * 100)
        
        for pub in active_pubs:
            name = pub['name']
            port = pub['port']
            host_str = pub['host']
            status = pub['status']
            
            if isinstance(pub['data_types'], dict):
                types_str = ', '.join([f"{k}:{v}" for k, v in pub['data_types'].items()])
            else:
                types_str = pub['data_types']
            
            status_display = '✓ Active' if status == 'active' else '✗ Inactive'
            print(f"{name:<20} {port:<8} {host_str:<15} {status_display:<10} {types_str}")
    
    def monitor_publisher(self, name_or_port, host='localhost'):
        if name_or_port.isdigit():
            port = int(name_or_port)
            addr = f"tcp://{host}:{port}"
        else:
            info = self.db.get_publisher(name_or_port)
            if not info:
                print(f"Publisher '{name_or_port}' not found in database")
                print("Use 'list' to see registered publishers or 'add' to register a new one")
                return
            port = info['port']
            host = info['host']
            addr = f"tcp://{host}:{port}"
        
        socket = self.context.socket(zmq.SUB)
        socket.setsockopt(zmq.SUBSCRIBE, b'')
        socket.connect(addr)
        
        print(f"Monitoring {addr}")
        print("Press Ctrl+C to stop\n")
        
        try:
            while True:
                try:
                    data = socket.recv_json(flags=zmq.NOBLOCK)
                    timestamp = datetime.now().strftime('%H:%M:%S.%f')[:-3]
                    
                    if isinstance(data, dict):
                        print(f"[{timestamp}]")
                        for key, value in data.items():
                            if isinstance(value, float):
                                print(f"  {key}: {value:.4f}")
                            else:
                                print(f"  {key}: {value}")
                        print()
                    else:
                        print(f"[{timestamp}] {data}\n")
                    
                except zmq.Again:
                    time.sleep(0.01)
                    
        except KeyboardInterrupt:
            print("\nStopped monitoring.")
        finally:
            socket.close()
    
    def add_publisher(self, name, port, host='localhost', description=''):
        if self.db.add_publisher(name, port, host, description):
            print(f"Added publisher '{name}' -> {host}:{port}")
        else:
            print(f"Publisher '{name}' already exists. Use 'update' to modify.")
    
    def remove_publisher(self, name):
        if self.db.remove_publisher(name):
            print(f"Removed publisher '{name}'")
        else:
            print(f"Publisher '{name}' not found")
    
    def list_registry(self):
        publishers = self.db.list_all_publishers()
        if not publishers:
            print("No registered publishers.")
            return
        
        print(f"{'Name':<20} {'Port':<8} {'Host':<15} {'Description'}")
        print("-" * 80)
        for pub in publishers:
            name = pub['name']
            port = pub['port']
            host = pub['host']
            desc = pub.get('description', '')
            print(f"{name:<20} {port:<8} {host:<15} {desc}")
    
    def show_config(self):
        start, end = self.db.get_port_range()
        count = end - start
        print(f"Port range configuration:")
        print(f"  Start port: {start}")
        print(f"  Port count: {count}")
        print(f"  End port:   {end - 1} (inclusive)")
        print(f"\nUsed ports: {len(self.db.get_used_ports())}/{count}")
    
    def set_port_range(self, start: int, count: int):
        if count < 10:
            print("Error: Port count must be at least 10")
            return
        if count > 10000:
            print("Error: Port count cannot exceed 10000")
            return
        
        self.db.set_port_range(start, count)
        print(f"Port range updated: {start} to {start + count - 1} ({count} ports)")
        print(f"Note: Existing publishers are not affected")
    
    def show_graph(self):
        graph = self.db.get_graph()
        
        if not graph:
            print("No publishers registered.")
            return
        
        print("Publisher → Subscriber Graph:\n")
        print("=" * 80)
        
        has_subscribers = False
        for pub_name, pub_info in graph.items():
            subscribers = pub_info['subscribers']
            port = pub_info['port']
            host = pub_info['host']
            
            if subscribers:
                has_subscribers = True
                print(f"\n📡 {pub_name} ({host}:{port})")
                for sub in subscribers:
                    print(f"   └─> {sub['name']} @ {sub['host']}")
            else:
                print(f"\n📡 {pub_name} ({host}:{port})")
                print(f"   (no subscribers)")
        
        print("\n" + "=" * 80)
        
        total_subs = sum(len(pub_info['subscribers']) for pub_info in graph.values())
        print(f"\nTotal: {len(graph)} publisher(s), {total_subs} subscription(s)")
    
    def show_subscriptions(self):
        subs = self.db.get_all_subscriptions()
        
        if not subs:
            print("No active subscriptions.")
            return
        
        print(f"{'Publisher':<20} {'Subscriber':<30} {'Host':<15} {'Last Active'}")
        print("-" * 90)
        
        for sub in subs:
            pub_name = sub['publisher_name']
            sub_name = sub['subscriber_name']
            sub_host = sub['subscriber_host']
            last_active = sub['last_active'].split('.')[0] if '.' in sub['last_active'] else sub['last_active']
            
            print(f"{pub_name:<20} {sub_name:<30} {sub_host:<15} {last_active}")


def main():
    parser = argparse.ArgumentParser(description='ZMQ Publisher Monitor')
    subparsers = parser.add_subparsers(dest='command', help='Commands')
    
    list_parser = subparsers.add_parser('list', help='List registered publishers and check their status')
    list_parser.add_argument('--host', default='localhost', help='Host to check (default: localhost)')
    
    reg_parser = subparsers.add_parser('registry', help='Show all registered publishers')
    
    graph_parser = subparsers.add_parser('graph', help='Show publisher-subscriber relationship graph')
    
    subs_parser = subparsers.add_parser('subscriptions', help='Show all active subscriptions')
    
    config_parser = subparsers.add_parser('config', help='Show or set port range configuration')
    config_parser.add_argument('--start', type=int, help='Start port number')
    config_parser.add_argument('--count', type=int, help='Number of ports in range')
    
    add_parser = subparsers.add_parser('add', help='Add a publisher to database')
    add_parser.add_argument('name', help='Publisher name')
    add_parser.add_argument('port', type=int, help='Port number')
    add_parser.add_argument('--host', default='localhost', help='Host address (default: localhost)')
    add_parser.add_argument('--desc', default='', help='Description')
    
    remove_parser = subparsers.add_parser('remove', help='Remove a publisher from database')
    remove_parser.add_argument('name', help='Publisher name to remove')
    
    monitor_parser = subparsers.add_parser('monitor', help='Monitor a publisher by name or port')
    monitor_parser.add_argument('name_or_port', help='Publisher name or port number')
    monitor_parser.add_argument('--host', default='localhost', help='Host (only used if port number given)')
    
    args = parser.parse_args()
    
    monitor = ZMQMonitor()
    
    if args.command == 'list':
        monitor.discover_and_list(host=args.host)
    elif args.command == 'registry':
        monitor.list_registry()
    elif args.command == 'graph':
        monitor.show_graph()
    elif args.command == 'subscriptions':
        monitor.show_subscriptions()
    elif args.command == 'config':
        if args.start is not None and args.count is not None:
            monitor.set_port_range(args.start, args.count)
        elif args.start is not None or args.count is not None:
            print("Error: Both --start and --count must be provided together")
        else:
            monitor.show_config()
    elif args.command == 'add':
        monitor.add_publisher(args.name, args.port, args.host, args.desc)
    elif args.command == 'remove':
        monitor.remove_publisher(args.name)
    elif args.command == 'monitor':
        monitor.monitor_publisher(args.name_or_port, getattr(args, 'host', 'localhost'))
    else:
        parser.print_help()


if __name__ == "__main__":
    main()