import argparse
import sys
import os
import logging

# Aggiungi le directory al path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def main():
    parser = argparse.ArgumentParser(description='Distributed LLM System')
    parser.add_argument('--mode', choices=['coordinator', 'worker', 'client'], required=True)
    parser.add_argument('--host', default='0.0.0.0')
    parser.add_argument('--port', type=int, default=8765)
    parser.add_argument('--coordinator-host', default='0.0.0.0')
    parser.add_argument('--coordinator-port', type=int, default=8765)
    parser.add_argument('--model-path', help='Path to LLM model')
    parser.add_argument('--gpu-id', type=int, default=0)
    
    args = parser.parse_args()
    
    try:
        if args.mode == 'coordinator':
            from coordinator.coordinator import Coordinator
            from blockchain.blockchain import Blockchain
            start_coordinator(args, Coordinator, Blockchain)
        elif args.mode == 'worker':
            from worker.gpu_worker import GPUWorker
            start_worker(args, GPUWorker)
        elif args.mode == 'client':
            from client.chat_client import ChatClient
            start_client(args, ChatClient)
    except ImportError as e:
        logger.error(f"Import error: {e}")
        logger.error("Make sure all required modules are available")
        sys.exit(1)
    except Exception as e:
        logger.error(f"Error: {e}")
        sys.exit(1)

def start_coordinator(args, Coordinator, Blockchain):
    """Avvia il coordinator"""
    blockchain = Blockchain()
    coordinator = Coordinator(args.host, args.port, blockchain)
    
    try:
        logger.info(f"Starting coordinator on {args.host}:{args.port}")
        coordinator.start()
    except KeyboardInterrupt:
        logger.info("Shutting down coordinator...")
        coordinator.stop()
    except Exception as e:
        logger.error(f"Coordinator error: {e}")

def start_worker(args, GPUWorker):
    """Avvia un worker GPU"""
    if not args.model_path:
        logger.error("Model path is required for worker mode")
        sys.exit(1)
        
    worker = GPUWorker(
        args.coordinator_host,
        args.coordinator_port,
        args.model_path,
        args.gpu_id
    )
    
    try:
        logger.info(f"Starting worker connecting to {args.coordinator_host}:{args.coordinator_port}")
        worker.connect()
    except KeyboardInterrupt:
        logger.info("Shutting down worker...")
        worker.disconnect()
    except Exception as e:
        logger.error(f"Worker error: {e}")

def start_client(args, ChatClient):
    """Avvia il client chat"""
    client = ChatClient(args.coordinator_host, args.coordinator_port)
    
    try:
        logger.info(f"Connecting to coordinator at {args.coordinator_host}:{args.coordinator_port}")
        client.start_chat()
    except KeyboardInterrupt:
        logger.info("Closing client...")
    except Exception as e:
        logger.error(f"Client error: {e}")

if __name__ == "__main__":
    main()