from __future__ import annotations
from jaclang import *
import typing
from enum import Enum, auto
if typing.TYPE_CHECKING:
    import io
else:
    io, = jac_import('io', 'py')
if typing.TYPE_CHECKING:
    import os
else:
    os, = jac_import('os', 'py')
if typing.TYPE_CHECKING:
    import re
else:
    re, = jac_import('re', 'py')
if typing.TYPE_CHECKING:
    import subprocess
else:
    subprocess, = jac_import('subprocess', 'py')
if typing.TYPE_CHECKING:
    import pkg_resources
else:
    pkg_resources, = jac_import('pkg_resources', 'py')
if typing.TYPE_CHECKING:
    import sys
else:
    sys, = jac_import('sys', 'py')
if typing.TYPE_CHECKING:
    import yaml
else:
    yaml, = jac_import('yaml', 'py')
if typing.TYPE_CHECKING:
    import logging
else:
    logging, = jac_import('logging', 'py')
if typing.TYPE_CHECKING:
    import traceback
else:
    traceback, = jac_import('traceback', 'py')
if typing.TYPE_CHECKING:
    import tarfile
else:
    tarfile, = jac_import('tarfile', 'py')
if typing.TYPE_CHECKING:
    import requests
else:
    requests, = jac_import('requests', 'py')
if typing.TYPE_CHECKING:
    from logging import Logger
else:
    Logger, = jac_import('logging', 'py', items={'Logger': None})
if typing.TYPE_CHECKING:
    from jaclang import jac_import
else:
    jac_import, = jac_import('jaclang', 'py', items={'jac_import': None})
if typing.TYPE_CHECKING:
    from jivas.agent.modules.agentlib.utils import Utils
else:
    Utils, = jac_import('jivas.agent.modules.agentlib.utils', 'py', items={'Utils': None})
if typing.TYPE_CHECKING:
    from jvcli.api import RegistryAPI
else:
    RegistryAPI, = jac_import('jvcli.api', 'py', items={'RegistryAPI': None})
if typing.TYPE_CHECKING:
    from jvcli.utils import is_version_compatible
else:
    is_version_compatible, = jac_import('jvcli.utils', 'py', items={'is_version_compatible': None})
if typing.TYPE_CHECKING:
    from jvserve.lib.agent_interface import AgentInterface
else:
    AgentInterface, = jac_import('jvserve.lib.agent_interface', 'py', items={'AgentInterface': None})
if typing.TYPE_CHECKING:
    from action import Action
else:
    Action, = jac_import('action', items={'Action': None})
if typing.TYPE_CHECKING:
    from interact_action import InteractAction
else:
    InteractAction, = jac_import('interact_action', items={'InteractAction': None})
if typing.TYPE_CHECKING:
    from jivas.agent.core.graph_node import GraphNode
else:
    GraphNode, = jac_import('jivas.agent.core.graph_node', items={'GraphNode': None})
if typing.TYPE_CHECKING:
    from exit_interact_action import ExitInteractAction
else:
    ExitInteractAction, = jac_import('exit_interact_action', items={'ExitInteractAction': None})
if typing.TYPE_CHECKING:
    from jivas.agent.core.purge import purge
else:
    purge, = jac_import('jivas.agent.core.purge', items={'purge': None})

class Actions(GraphNode, Node):
    logger: static[Logger] = logging.getLogger(__name__)
    package_root: static[str] = Utils.get_actions_root()

    def get(self, action_type: str='', action_label: str='', only_enabled: bool=True) -> None:
        result = None
        if action_type and (not action_label):
            result = self.get_by_type(action_type, only_enabled)
        elif action_label and (not action_type):
            result = self.get_by_label(action_label, only_enabled)
        return result

    def get_by_label(self, action_label: str, only_enabled: bool=True) -> Action:
        actions = self.get_all(only_enabled=only_enabled)
        for action_node in actions:
            if action_node.label == action_label:
                return action_node
        return None

    def get_by_type(self, action_type: str, only_enabled: bool=True) -> list:
        result = JacList([])
        actions = self.get_all(only_enabled=only_enabled)
        for action_node in actions:
            if action_type == action_node.get_type():
                result.append(action_node)
        return result

    def get_all(self, only_interact_actions: bool=False, only_enabled: bool=False) -> list:
        if only_interact_actions:
            return self.spawn(_get_interact_actions(filter_enabled=only_enabled)).action_nodes
        return self.spawn(_get_actions(filter_enabled=only_enabled)).action_nodes

    def queue_interact_actions(self, actions: list) -> list:
        return sorted(actions, key=lambda action: InteractAction: action.weight)

    def get_action_info(self, namespace_package_name: str, version: str=None, jpr_api_key: str=None) -> dict:
        action_info = {}
        namespace, package_name = namespace_package_name.split('/')
        if (package_path := Utils.find_package_folder(self.package_root, namespace_package_name)):
            info_yaml_path = os.path.join(package_path, 'info.yaml')
            try:
                with open(info_yaml_path, 'r') as file:
                    _info = yaml.safe_load(file)
                    package_version = _info.get('package', {}).get('version', '>=0.0.1')
                    if is_version_compatible(package_version, version):
                        self.logger.info(f'{namespace_package_name} {package_version} loaded locally')
                        module_root = Utils.path_to_module(package_path)
                        has_app = os.path.isfile(os.path.join(package_path, 'app', 'app.py'))
                        action_info = _info['package']
                        action_info['config']['path'] = package_path
                        action_info['config']['app'] = has_app
                        action_info['config']['module_root'] = module_root
                        action_info['config']['module'] = f'{module_root}.{package_name}'
            except yaml.YAMLError as e:
                self.logger.error(f'an exception occurred, {traceback.format_exc()}')
        if not action_info:
            if (_info := RegistryAPI.get_package_info(namespace_package_name, version, api_key=jpr_api_key)):
                self.logger.info(f'{namespace_package_name} {version} found in registry')
                action_info = _info.get('package', None)
        if action_info:
            action_info['config']['namespace'] = namespace
            action_info['config']['package_name'] = package_name
        return action_info

    def import_action(self, action_data: dict) -> bool:
        module_root = action_data.get('context', {}).get('_package', {}).get('config', {}).get('module_root', None)
        if module_root:
            jac_import(f'{module_root}.lib', base_path='./')
            return True
        return False

    def register_action(self, action_data: dict, parent: str='') -> Action:
        action_node = None
        if not action_data:
            self.logger.error(f'unable to register action {label}, missing or invalid action data')
            return None
        label = action_data.get('context', {}).get('label', action_data.get('action', None))
        architype = action_data.get('context', {}).get('_package', {}).get('architype', None)
        module = action_data.get('context', {}).get('_package', {}).get('config', {}).get('module', None)
        singleton = action_data.get('context', {}).get('_package', {}).get('config', {}).get('singleton', False)
        action_type = action_data.get('context', {}).get('_package', {}).get('meta', {}).get('type', 'action')
        if not architype or not module or (not label):
            self.logger.error(f'unable to register action {label}, missing label, architype or module name')
            return None
        try:
            if singleton:
                if (existing_action := self.get(action_type=label, only_enabled=False)):
                    self.logger.error(f'action already exists: {label}')
                    return None
            elif (existing_action := self.get(action_label=label, only_enabled=False)):
                self.logger.error(f'action already exists: {existing_action.label}')
                return None
            action_node = AgentInterface.spawn_node(node_name=architype, attributes=action_data.get('context', {}), module_name=module)
            if action_node:
                if not parent:
                    action_parent_node = self
                else:
                    action_parent_node = self.get_by_type(action_type=parent, only_enabled=False)
                action_parent_node.connect(action_node)
                action_node.on_register()
                self.logger.info(f'registered action: {action_node.label}')
            if 'children' in action_data and action_type == 'interact_action':
                for child_data in action_data['children']:
                    self.register_action(action_data=child_data, parent=architype)
        except Exception as e:
            self.logger.error(f'an exception occurred wile registering action {label}, {traceback.format_exc()}')

    def install_actions(self, agent_id: str, data: dict, jpr_api_key: str=None) -> bool:
        loaded_actions_data = None
        if agent_id:
            loaded_actions_data = self.load_action_packages(agent_id, data, jpr_api_key=jpr_api_key)
        if not loaded_actions_data:
            self.logger.error('no actions loaded; unable to proceed with import')
            return False
        self.deregister_actions()
        for action_data in loaded_actions_data:
            self.register_action(action_data=action_data)
        self.connect(ExitInteractAction())
        for action_node in self.get_all():
            action_node.post_register()
        return True

    def deregister_action(self, action_type: str='', action_label: str='') -> None:
        target = JacList([])
        if action_type and (not action_label):
            target = self.get_by_type(action_type=action_type, only_enabled=False)
        elif action_label and (not action_type):
            target.append(self.get_by_label(action_label=action_label, only_enabled=False))
        for action_node in target:
            action_node.on_deregister()
            action_node.spawn(purge())

    def deregister_actions(self) -> None:
        for action_node in self.get_all():
            action_node.on_deregister()
        self.spawn(purge(purge_spawn_node=False))
        Utils.jac_clean_actions()

    def index_action_packages(self, agent_id: str, action_list: list, action_index: dict={}, jpr_api_key: str='') -> dict:
        for action_data in action_list:
            is_not_indexed = action_data.get('context', {}).get('_package', None) is None
            namespace_package_name = action_data.get('action', None)
            package_version = action_data.get('context', {}).get('version', '>=0.0.1')
            if is_not_indexed and namespace_package_name:
                if (action_info := self.get_action_info(namespace_package_name, package_version, jpr_api_key)):
                    action_data['context']['agent_id'] = agent_id
                    action_data['context']['_package'] = action_info
                    architype = action_info.get('architype', None)
                    action_data['context']['label'] = action_data.get('context', {}).get('label', architype)
                    action_data['context']['description'] = action_info.get('meta', {}).get('description', '')
                    if (action_deps := action_info.get('dependencies', {}).get('actions', {})):
                        action_dep_list = JacList([])
                        for action_dep in action_deps.keys():
                            action_dep_version = action_deps[action_dep]
                            action_dep_list.append({'action': action_dep, 'context': {'version': action_dep_version, 'enabled': True}})
                        action_index = self.index_action_packages(agent_id=agent_id, action_list=action_dep_list, action_index=action_index, jpr_api_key=jpr_api_key)
                        if not all((key in action_index.keys() for key in action_deps.keys())):
                            continue
                    if not action_index.get(namespace_package_name):
                        action_index[namespace_package_name] = action_data
                else:
                    self.logger.error(f'unable to find action {namespace_package_name} {package_version}')
                if 'children' in action_data:
                    action_index = self.index_action_packages(agent_id=agent_id, action_list=action_data['children'], action_index=action_index, jpr_api_key=jpr_api_key)
        return action_index

    def index_pip_packages(self, action_data_index: dict) -> dict:
        pip_packages_index = {}
        for namespace_package_name, action_data in action_data_index.items():
            if (pip_packages := action_data.get('context', {}).get('_package', {}).get('dependencies', {}).get('pip', {})):
                pip_packages_index.update(pip_packages)
        return pip_packages_index

    def batch_pip_package_install(self, packages: dict) -> None:
        package_specs = JacList([])
        for pkg, ver in packages.items():
            if not ver:
                package_specs.append(pkg)
            elif bool(re.match('^(==|>=|<=|>|<|~=|!=)', ver)):
                package_specs.append(f'{pkg}{ver}')
            else:
                package_specs.append(f'{pkg}=={ver}')
        command = JacList([sys.executable, '-m', 'pip', 'install', '--upgrade', '--no-input'])
        command.extend(package_specs)
        self.logger.info(f'Attempting to install package dependencies: {package_specs}')
        try:
            subprocess.run(command, check=True, stderr=None, stdout=None)
            self.logger.info('Package dependencies installed successfully.')
        except subprocess.CalledProcessError as e:
            self.logger.error('An error occurred while installing one or more package dependencies. There may be a conflict or a missing package.')

    def has_pip_package(self, package_name: str, version: str=None) -> bool:
        try:
            package_distribution = pkg_resources.get_distribution(package_name)
            if version is None:
                return True
            else:
                if not re.match('^(==|>=|<=|>|<|~=|!=)', version):
                    version = f'=={version}'
                return pkg_resources.parse_version(package_distribution.version) in pkg_resources.Requirement.parse(f'{package_name}{version}')
        except pkg_resources.DistributionNotFound:
            return False

    def has_pip_packages(self, packages: dict) -> bool:
        for pkg, ver in packages.items():
            if not self.has_pip_package(pkg, ver):
                return False
        return True

    def has_action_dependencies(self, action_data_index: dict, action_dependencies: dict) -> bool:
        if not all((key in action_data_index.keys() for key in action_dependencies.keys())):
            return False
        return True

    def load_action_package(self, indexed_action_data: dict, jpr_api_key: str='') -> dict:
        namespace_package_name = indexed_action_data.get('context', {}).get('_package', {}).get('name')
        package_version = indexed_action_data.get('context', {}).get('version', '>=0.0.1')
        jivas_version = indexed_action_data.get('context', {}).get('_package', {}).get('dependencies', {}).get('jivas', None)
        package_path = indexed_action_data.get('context', {}).get('_package', {}).get('config', {}).get('path', None)
        if not is_version_compatible('2.0.0-alpha.35', '^2.0.0'):
            self.logger.error(f'incompatible JIVAS version for: {namespace_package_name} ( {jivas_version} : {Utils.get_jivas_version()} ) ...skipping')
            return {}
        if package_path is None:
            package_data = RegistryAPI.download_package(namespace_package_name, package_version, api_key=jpr_api_key)
            if not package_data:
                self.logger.error(f'unable to load action package: {namespace_package_name} {package_version} ...skipping')
                return {}
            try:
                package_file = requests.get(package_data['file'])
                target_dir = os.path.join(self.package_root, namespace_package_name)
                os.makedirs(target_dir, exist_ok=True)
                with tarfile.open(fileobj=io.BytesIO(package_file.content), mode='r:gz') as tar_file:
                    tar_file.extractall(target_dir)
                self.logger.info(f'downloaded action package: {namespace_package_name} {package_version}')
            except Exception as e:
                self.logger.error(f'unable to save action package: {namespace_package_name}...skipping, {e}')
                return {}
            if (action_info := self.get_action_info(namespace_package_name, package_version, jpr_api_key)):
                indexed_action_data['context']['_package'] = action_info
                if not action_info.get('config', {}).get('path', None):
                    self.logger.error(f'unable to load action package: {namespace_package_name} {package_version}...skipping')
                    return {}
        return indexed_action_data

    def load_action_packages(self, agent_id: str, data: list, jpr_api_key: str=None) -> list:
        action_data_index = self.index_action_packages(agent_id=agent_id, action_list=data, jpr_api_key=jpr_api_key)
        if not action_data_index:
            return JacList([])
        subprocess.check_call(JacList([sys.executable, '-m', 'pip', 'install', '--upgrade', 'pip', '--root-user-action=ignore']))
        indexed_pip_packages = self.index_pip_packages(action_data_index)
        self.batch_pip_package_install(indexed_pip_packages)
        action_data_index_copy = action_data_index.copy()
        for namespace_package_name, action_data in action_data_index_copy.items():
            loaded_action_data = self.load_action_package(indexed_action_data=action_data, jpr_api_key=jpr_api_key)
            if not loaded_action_data:
                del action_data_index[namespace_package_name]
                continue
            if not self.has_pip_packages(loaded_action_data.get('context', {}).get('_package', {}).get('dependencies', {}).get('pip', {})):
                self.logger.error(f'missing pip package dependencies for {namespace_package_name}...skipping')
                del action_data_index[namespace_package_name]
                continue
            action_data_index[namespace_package_name] = loaded_action_data
            action_deps = loaded_action_data.get('context', {}).get('_package', {}).get('dependencies', {}).get('actions', {})
            if not self.has_action_dependencies(action_data_index, action_deps):
                self.logger.error(f'missing action dependencies for {namespace_package_name}...skipping')
                del action_data_index[namespace_package_name]
                continue
            self.import_action(loaded_action_data)
            if 'children' in loaded_action_data and action_type == 'interact_action':

                def update_child_actions(child_action_list):
                    for i, child_action_data in enumerate(child_action_list):
                        action_type = child_action_data.get('context', {}).get('_package', {}).get('meta', {}).get('type', 'action')
                        if (namespace_package_name := child_action_data.get('action', None)):
                            if namespace_package_name in action_data_index:
                                child_action_list[i] = action_data_index[namespace_package_name]
                                del action_data_index[namespace_package_name]
                            if 'children' in child_action_data and action_type == 'interact_action':
                                update_child_actions(child_action_data['children'])
                update_child_actions(loaded_action_data['children'])
        return Utils.order_interact_actions(action_data_index.values())

class _get_actions(Walker):
    action_nodes: list = field(gen=lambda: JacList([]))
    filter_enabled: bool = field(False)

    class __specs__(Obj):
        private: static[bool] = True

    @with_entry
    def on_actions(self, here: Actions) -> None:
        if self.filter_enabled:
            self.visit(here.refs().filter(Action, None).filter(None, lambda item: item.enabled == True))
        else:
            self.visit(here.refs().filter(Action, None))

    @with_entry
    def on_action(self, here: Action) -> None:
        self.action_nodes.append(here)
        self.visit(here.refs().filter(Action, None))

class _get_interact_actions(Walker):
    action_nodes: list = field(gen=lambda: JacList([]))
    filter_enabled: bool = field(False)

    class __specs__(Obj):
        private: static[bool] = True

    @with_entry
    def on_actions(self, here: Actions) -> None:
        if self.filter_enabled:
            self.visit(here.refs().filter(InteractAction, None).filter(None, lambda item: item.enabled == True))
        else:
            self.visit(here.refs().filter(InteractAction, None))

    @with_entry
    def on_action(self, here: InteractAction) -> None:
        self.action_nodes.append(here)
        self.visit(here.refs().filter(InteractAction, None))