import shlex
from .azure_connector import AzureConnector

from .azure_agents.vm_agent import VMAgent
from .azure_agents.storage_agent import StorageAgent
from .azure_agents.sql_agent import SQLAgent
from .azure_agents.functions_agent import FunctionsAgent

class AzureAgentManager:
    """
    The central orchestrator for all Azure agents.
    """

    def __init__(self):
        """
        Initializes the AzureAgentManager and registers the available agents.
        """
        self.agents = {
            'vm': VMAgent(),
            'storage': StorageAgent(),
            'sql': SQLAgent(),
            'functions': FunctionsAgent(),
        }
        print("AzureAgentManager initialized with the following agents: " + ", ".join(self.agents.keys()))

    def get_available_services(self):
        """
        Returns a list of all services managed by this AzureAgentManager.
        """
        return list(self.agents.keys())

    def get_service_commands(self, service_name: str):
        """
        Returns a list of commands supported by a specific service agent.
        """
        parser_method = getattr(self, f'_parse_{service_name}_command', None)
        if parser_method:
            # Call the parser with dummy parts to get the command list
            # The parser now returns (command, args, commands_list)
            _, _, commands_list = parser_method([])
            if commands_list:
                return [cmd['usage'] for cmd in commands_list]
        return []

    def _get_service(self, parts):
        # Implement logic to determine Azure service from command parts
        if 'vm' in parts or 'virtual-machine' in parts:
            return 'vm'
        if 'storage' in parts or 'storage-account' in parts:
            return 'storage'
        if 'sql' in parts or 'database' in parts:
            return 'sql'
        if 'function' in parts or 'functions' in parts:
            return 'functions'
        if 'vnet' in parts or 'virtual-network' in parts:
            return 'network'
        if 'aks' in parts or 'kubernetes' in parts or 'cluster' in parts:
            return 'aks'
        return None

    # Placeholder for Azure service command parsers
    def _parse_vm_command(self, parts):
        commands = [
            {"command": "list_vms", "usage": "vm list vms resource-group <name> in <location>"},
            {"command": "start_vm", "usage": "vm start <vm-name> resource-group <name> in <location>"},
            {"command": "stop_vm", "usage": "vm stop <vm-name> resource-group <name> in <location>"},
            {"command": "delete_vm", "usage": "vm delete <vm-name> resource-group <name> in <location>"},
            {"command": "smart_create_vm", "usage": "vm smart-create vm <name> resource-group <rg> in <location> [image <publisher> <offer> <sku>] [size <vm-size>] [username <user>] [password <pass>]"},
        ]

        command, args = None, {}
        try:
            args['subscription_id'] = AzureConnector.get_subscription_id()
            if 'resource-group' in parts:
                args['resource_group'] = parts[parts.index('resource-group') + 1]
            else:
                return None, {"error": "Resource group not specified. Use 'resource-group <name>'."}, commands
            if 'in' in parts:
                args['location'] = parts[parts.index('in') + 1]
            else:
                return None, {"error": "Location not specified. Use 'in <location>'."}, commands

            if 'list' in parts and ('vms' in parts or 'virtual-machines' in parts):
                command = 'list_vms'
            elif 'start' in parts or 'stop' in parts or 'delete' in parts:
                action = parts[0]
                command = {'start': 'start_vm', 'stop': 'stop_vm', 'delete': 'delete_vm'}[action]
                try:
                    args['vm_name'] = parts[parts.index(action) + 1]
                except (ValueError, IndexError):
                    return None, {"error": f"VM name not specified for '{action}'."}, commands
            elif 'smart-create' in parts and ('vm' in parts or 'virtual-machine' in parts):
                command = 'smart_create_vm'
                try:
                    args['vm_name'] = parts[parts.index('vm') + 1]
                    if 'image' in parts:
                        args['image_publisher'] = parts[parts.index('image') + 1]
                        args['image_offer'] = parts[parts.index('image') + 2]
                        args['image_sku'] = parts[parts.index('image') + 3]
                    if 'size' in parts:
                        args['vm_size'] = parts[parts.index('size') + 1]
                    if 'username' in parts:
                        args['admin_username'] = parts[parts.index('username') + 1]
                    if 'password' in parts:
                        args['admin_password'] = parts[parts.index('password') + 1]
                except (ValueError, IndexError):
                    return None, {"error": "Usage: vm smart-create vm <name> resource-group <rg> in <location> [image <publisher> <offer> <sku>] [size <vm-size>] [username <user>] [password <pass>]"}, commands
        except (ValueError, IndexError) as e:
            return None, {"error": f"Invalid VM command usage: {e}. Refer to help for correct syntax."}
        return command, args, commands

    def _parse_storage_command(self, parts):
        commands = [
            {"command": "list_storage_accounts", "usage": "storage list storage-accounts resource-group <name> in <location>"},
            {"command": "smart_create_storage_account", "usage": "storage smart-create storage-account <name> resource-group <rg> in <location>"},
            {"command": "delete_storage_account", "usage": "storage delete storage-account <name> resource-group <rg>"},
        ]

        command, args = None, {}
        try:
            args['subscription_id'] = AzureConnector.get_subscription_id()
            if 'resource-group' in parts:
                args['resource_group'] = parts[parts.index('resource-group') + 1]
            else:
                return None, {"error": "Resource group not specified. Use 'resource-group <name>'."}, commands
            if 'in' in parts:
                args['location'] = parts[parts.index('in') + 1]
            else:
                return None, {"error": "Location not specified. Use 'in <location>'."}, commands

            if 'list' in parts and 'storage-accounts' in parts:
                command = 'list_storage_accounts'
            elif 'smart-create' in parts and 'storage-account' in parts:
                command = 'smart_create_storage_account'
                try:
                    args['account_name'] = parts[parts.index('storage-account') + 1]
                except (ValueError, IndexError):
                    return None, {"error": "Usage: storage smart-create storage-account <name> resource-group <rg> in <location>"}, commands
            elif 'delete' in parts and 'storage-account' in parts:
                command = 'delete_storage_account'
                try:
                    args['account_name'] = parts[parts.index('storage-account') + 1]
                except (ValueError, IndexError):
                    return None, {"error": "Usage: storage delete storage-account <name> resource-group <rg>"}, commands
        except (ValueError, IndexError) as e:
            return None, {"error": f"Invalid Storage command usage: {e}. Refer to help for correct syntax."}
        return command, args, commands

    def _parse_sql_command(self, parts):
        commands = [
            {"command": "smart_create_sql_database", "usage": "sql smart-create sql-database server <server-name> database <db-name> username <user> password <pass> resource-group <rg> in <location>"},
        ]

        command, args = None, {}
        try:
            args['subscription_id'] = AzureConnector.get_subscription_id()
            if 'resource-group' in parts:
                args['resource_group'] = parts[parts.index('resource-group') + 1]
            else:
                return None, {"error": "Resource group not specified. Use 'resource-group <name>'."}, commands
            if 'in' in parts:
                args['location'] = parts[parts.index('in') + 1]
            else:
                return None, {"error": "Location not specified. Use 'in <location>'."}, commands

            if 'smart-create' in parts and 'sql-database' in parts:
                command = 'smart_create_sql_database'
                try:
                    args['server_name'] = parts[parts.index('server') + 1]
                    args['database_name'] = parts[parts.index('database') + 1]
                    args['admin_username'] = parts[parts.index('username') + 1]
                    args['admin_password'] = parts[parts.index('password') + 1]
                except (ValueError, IndexError):
                    return None, {"error": "Usage: sql smart-create sql-database server <server-name> database <db-name> username <user> password <pass> resource-group <rg> in <location>"}, commands
        except (ValueError, IndexError) as e:
            return None, {"error": f"Invalid SQL command usage: {e}. Refer to help for correct syntax."}
        return command, args, commands

    def _parse_functions_command(self, parts):
        commands = [
            {"command": "smart_create_function_app", "usage": "function smart-create function <app-name> storage-account <sa-name> resource-group <rg> in <location> [runtime <runtime>] [runtime-version <version>]"},
        ]

        command, args = None, {}
        try:
            args['subscription_id'] = AzureConnector.get_subscription_id()
            if 'resource-group' in parts:
                args['resource_group'] = parts[parts.index('resource-group') + 1]
            else:
                return None, {"error": "Resource group not specified. Use 'resource-group <name>'."}, commands
            if 'in' in parts:
                args['location'] = parts[parts.index('in') + 1]
            else:
                return None, {"error": "Location not specified. Use 'in <location>'."}, commands

            if 'smart-create' in parts and 'function' in parts:
                command = 'smart_create_function_app'
                try:
                    args['app_name'] = parts[parts.index('function') + 1]
                    args['storage_account_name'] = parts[parts.index('storage-account') + 1]
                    if 'runtime' in parts:
                        args['runtime'] = parts[parts.index('runtime') + 1]
                    if 'runtime-version' in parts:
                        args['runtime_version'] = parts[parts.index('runtime-version') + 1]
                except (ValueError, IndexError):
                    return None, {"error": "Usage: function smart-create function <app-name> storage-account <sa-name> resource-group <rg> in <location> [runtime <runtime>] [runtime-version <version>]"}, commands
        except (ValueError, IndexError) as e:
            return None, {"error": f"Invalid Function command usage: {e}. Refer to help for correct syntax."}
        return command, args, commands

    def _parse_network_command(self, parts):
        commands = [
            {"command": "smart_create_vnet", "usage": "vnet smart-create vnet <name> resource-group <rg> in <location> [address-prefix <cidr>] [subnet-prefix <cidr>]"},
        ]

        command, args = None, {}
        try:
            args['subscription_id'] = AzureConnector.get_subscription_id()
            if 'resource-group' in parts:
                args['resource_group'] = parts[parts.index('resource-group') + 1]
            else:
                return None, {"error": "Resource group not specified. Use 'resource-group <name>'."}, commands
            if 'in' in parts:
                args['location'] = parts[parts.index('in') + 1]
            else:
                return None, {"error": "Location not specified. Use 'in <location>'."}, commands

            if 'smart-create' in parts and ('vnet' in parts or 'virtual-network' in parts):
                command = 'smart_create_vnet'
                try:
                    args['vnet_name'] = parts[parts.index('vnet') + 1]
                    if 'address-prefix' in parts:
                        args['address_prefix'] = parts[parts.index('address-prefix') + 1]
                    if 'subnet-prefix' in parts:
                        args['subnet_prefix'] = parts[parts.index('subnet-prefix') + 1]
                except (ValueError, IndexError):
                    return None, {"error": "Usage: vnet smart-create vnet <name> resource-group <rg> in <location> [address-prefix <cidr>] [subnet-prefix <cidr>]"}, commands
        except (ValueError, IndexError) as e:
            return None, {"error": f"Invalid Network command usage: {e}. Refer to help for correct syntax."}
        return command, args, commands

    def _parse_aks_command(self, parts):
        commands = [
            {"command": "smart_create_aks_cluster", "usage": "aks smart-create aks-cluster <name> resource-group <rg> in <location> [node-count <count>] [node-vm-size <size>]"},
        ]

        command, args = None, {}
        try:
            args['subscription_id'] = AzureConnector.get_subscription_id()
            if 'resource-group' in parts:
                args['resource_group'] = parts[parts.index('resource-group') + 1]
            else:
                return None, {"error": "Resource group not specified. Use 'resource-group <name>'."}, commands
            if 'in' in parts:
                args['location'] = parts[parts.index('in') + 1]
            else:
                return None, {"error": "Location not specified. Use 'in <location>'."}, commands

            if 'smart-create' in parts and ('aks-cluster' in parts or 'cluster' in parts):
                command = 'smart_create_aks_cluster'
                try:
                    args['cluster_name'] = parts[parts.index('cluster') + 1]
                    if 'node-count' in parts:
                        args['node_count'] = int(parts[parts.index('node-count') + 1])
                    if 'node-vm-size' in parts:
                        args['node_vm_size'] = parts[parts.index('node-vm-size') + 1]
                except (ValueError, IndexError):
                    return None, {"error": "Usage: aks smart-create aks-cluster <name> resource-group <rg> in <location> [node-count <count>] [node-vm-size <size>]"}, commands
        except (ValueError, IndexError) as e:
            return None, {"error": f"Invalid AKS command usage: {e}. Refer to help for correct syntax."}
        return command, args, commands

    def _parse_command(self, user_command: str):
        parts = shlex.split(user_command)
        if not parts:
            return None, None, None

        # Extract subscription ID if provided as a global argument
        subscription_id = None
        if '--subscription' in parts:
            try:
                sub_index = parts.index('--subscription')
                subscription_id = parts[sub_index + 1]
                AzureConnector.set_subscription_id(subscription_id)
                # Remove from parts so it doesn't interfere with service parsing
                parts.pop(sub_index)
                parts.pop(sub_index) # Pop the value too
            except (ValueError, IndexError):
                return None, None, {"error": "Usage: --subscription <subscription-id>"}

        service = self._get_service(parts)
        if not service:
            return None, None, {}

        parser_map = {
            'vm': self._parse_vm_command,
            'storage': self._parse_storage_command,
            'sql': self._parse_sql_command,
            'functions': self._parse_functions_command,
            'network': self._parse_network_command,
            'aks': self._parse_aks_command,
        }

        parser = parser_map.get(service)
        if parser:
            command, args, commands_list = parser(parts)
            # Add subscription_id to args if it was parsed globally
            if subscription_id and args is not None:
                args['subscription_id'] = subscription_id
            return service, command, args, commands_list
        
        return None, None, {}, []

    def execute_command(self, user_command: str):
        """
        Parses and executes a user command, routing to Azure agents.
        """
        print(f"\nReceived Azure command: '{user_command}'")
        parts = shlex.split(user_command)

        service, command, args, commands_list = self._parse_command(user_command)

        if args and 'error' in args:
            return {"status": "error", "message": args['error']}

        if not service or not command:
            error_message = "Could not understand the Azure command. Please be more specific."
            if commands_list:
                error_message += "\nAvailable commands for this service:\n" + "\n".join([cmd['usage'] for cmd in commands_list])
            return {"status": "error", "message": error_message}

        if service in self.agents:
            agent = self.agents[service]
            try:
                print(f"Routing to {service.upper()}Agent with command '{command}' and args {args}")
                return agent.execute(command, **args)
            except Exception as e:
                return {"status": "error", "message": str(e)}
        else:
            return {"status": "error", "message": f"No agent found for Azure service '{service}'."}