import cmd
import socket
import json
import os

class ModuleShell(cmd.Cmd):
    intro = "Welcome to the Module Shell. Type 'help' for commands."
    prompt = "(module_shell) "
    
    def __init__(self, host='localhost', port=9999):
        super().__init__()
        self.host = host
        self.port = port
        self.update_commands()

    def send_command(self, command):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.connect((self.host, self.port))
            s.sendall(json.dumps(command).encode())
            return json.loads(s.recv(1024).decode())

    def update_commands(self):
        response = self.send_command({'action': 'list_commands'})
        if response['success']:
            self.dynamic_commands = response['message']
        else:
            self.dynamic_commands = []

    def default(self, line):
        parts = line.split()
        command = parts[0]
        args = ' '.join(parts[1:])
        response = self.send_command({'action': 'execute', 'command': command, 'args': args})
        if response['success']:
            print(response['message'])
        else:
            print(f"Error: {response['message']}")

    def completenames(self, text, *ignored):
        dotext = 'do_'+text
        return [a[3:] for a in self.get_names() if a.startswith(dotext)] + \
               [cmd for cmd in self.dynamic_commands if cmd.startswith(text)]

    def do_load(self, arg):
        """Load a module: load <module_name>"""
        response = self.send_command({'action': 'load', 'module': arg})
        print(response['message'])
        self.update_commands()  # Update commands after loading a new module

    def do_unload(self, arg):
        """Unload a module: unload <module_name>"""
        response = self.send_command({'action': 'unload', 'module': arg})
        print(response['message'])
        self.update_commands()  # Update commands after unloading a module

    def do_list(self, arg):
        """List all loaded modules"""
        response = self.send_command({'action': 'list'})
        modules = response['message']
        if modules:
            print("Loaded modules:")
            for module in modules:
                print(f"- {module}")
        else:
            print("No modules are currently loaded.")

    def do_list_available(self, arg):
        """List available modules in the modules folder"""
        modules_path = os.path.join(os.path.dirname(__file__), 'modules')
        available_modules = [f[:-3] for f in os.listdir(modules_path) if f.endswith('.py') and not f.startswith('__')]
        print("Available modules:")
        for module in available_modules:
            print(f"- {module}")

    def complete_load(self, text, line, begidx, endidx):
        modules_path = os.path.join(os.path.dirname(__file__), 'modules')
        available_modules = [f[:-3] for f in os.listdir(modules_path) if f.endswith('.py') and not f.startswith('__')]
        return [m for m in available_modules if m.startswith(text)]


    def do_upgrade(self, arg):
        """Upgrade the core daemon"""
        response = self.send_command({'action': 'upgrade'})
        print(response['message'])

    def do_exit(self, arg):
        """Exit the shell"""
        print("Exiting...")
        return True

def main():
    ModuleShell().cmdloop()

if __name__ == "__main__":
    main()