import os
import sys
import importlib
import socket
import json
import threading

# Add the modules folder to the Python path
modules_path = os.path.join(os.path.dirname(__file__), 'modules')
sys.path.append(modules_path)

class ModuleManager:
    def __init__(self, module_dirs):
        self.module_dirs = module_dirs
        self.loaded_modules = {}
        self.extra_commands = {}
        self._update_sys_path()

    def _update_sys_path(self):
        for dir in self.module_dirs:
            full_path = os.path.abspath(dir)
            if full_path not in sys.path:
                sys.path.append(full_path)

    def add_module_dir(self, new_dir):
        if new_dir not in self.module_dirs:
            self.module_dirs.append(new_dir)
            self._update_sys_path()
            return True, f"Added module directory: {new_dir}"
        return False, f"Module directory already exists: {new_dir}"

    def load_module(self, module_name):
        for dir in self.module_dirs:
            try:
                module = importlib.import_module(f'{dir}.{module_name}')
                self.loaded_modules[module_name] = module
                if hasattr(module, 'initialize'):
                    module.initialize()
                if hasattr(module, 'get_commands'):
                    new_commands = module.get_commands()
                    self.extra_commands.update(new_commands)
                return True, f"Module '{module_name}' loaded and initialized successfully from {dir}."
            except ImportError:
                continue
        return False, f"Error: Unable to load module '{module_name}' from any of the module directories."

    def unload_module(self, module_name):
        if module_name in self.loaded_modules:
            module = self.loaded_modules[module_name]
            if hasattr(module, 'shutdown'):
                module.shutdown()
            if hasattr(module, 'get_commands'):
                commands_to_remove = module.get_commands().keys()
                for cmd in commands_to_remove:
                    self.extra_commands.pop(cmd, None)
            del self.loaded_modules[module_name]
            return True, f"Module '{module_name}' unloaded and shut down."
        return False, f"Module '{module_name}' is not loaded."

    def list_modules(self):
        return list(self.loaded_modules.keys())

    def list_commands(self):
        return list(self.extra_commands.keys())

    def execute_command(self, command, args):
        if command in self.extra_commands:
            return True, self.extra_commands[command](args)
        return False, "Command not found"

class CoreDaemon:
    def __init__(self, host='localhost', port=9999, module_dirs=None):
        self.host = host
        self.port = port
        self.module_dirs = module_dirs or ['modules']  # Default to 'modules' if not specified
        self.module_manager = ModuleManager(self.module_dirs)

    def start(self):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind((self.host, self.port))
            s.listen()
            print(f"Core daemon listening on {self.host}:{self.port}")
            while True:
                conn, addr = s.accept()
                threading.Thread(target=self.handle_client, args=(conn,)).start()

    def handle_client(self, conn):
        with conn:
            while True:
                data = conn.recv(1024)
                if not data:
                    break
                command = json.loads(data.decode())
                response = self.process_command(command)
                conn.sendall(json.dumps(response).encode())

    def upgrade(self):
        print("Upgrading core daemon...")
        # Here you would typically download or copy the new version
        # For this example, we'll just restart the script
        os.execv(sys.executable, ['python'] + sys.argv)

    def process_command(self, command):
        action = command.get('action')
        if action == 'load':
            success, message = self.module_manager.load_module(command.get('module'))
        elif action == 'unload':
            success, message = self.module_manager.unload_module(command.get('module'))
        elif action == 'list':
            modules = self.module_manager.list_modules()
            success, message = True, modules
        elif action == 'execute':
            success, message = self.module_manager.execute_command(command.get('command'), command.get('args'))
        elif action == 'upgrade':
            self.upgrade()
            success, message = True, "Upgrade initiated"
        elif action == 'list_commands':
            commands = self.module_manager.list_commands()
            success, message = True, commands
        elif action == 'add_module_dir':
            success, message = self.module_manager.add_module_dir(command.get('dir'))
        else:
            success, message = False, "Unknown command"
        return {'success': success, 'message': message}

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Core Daemon with dynamic module loading")
    parser.add_argument('--host', default='localhost', help='Host to bind the daemon to')
    parser.add_argument('--port', type=int, default=9999, help='Port to bind the daemon to')
    parser.add_argument('--module-dirs', nargs='*', default=['modules'], help='Directories to load modules from')
    
    args = parser.parse_args()

    daemon = CoreDaemon(host=args.host, port=args.port, module_dirs=args.module_dirs)
    daemon.start()