import copy
import os
import pickle
import shutil
import tarfile
import uuid

from datetime import datetime

import app.service.data_svc
from app.utility.base_service import BaseService
from app.objects.secondclass.c_fact import Fact, WILDCARD_STRING
from app.objects.secondclass.c_relationship import Relationship

DATA_BACKUP_DIR = app.service.data_svc.DATA_BACKUP_DIR
FACT_STORE_PATH = f"data{os.path.sep}fact_store"


class BaseKnowledgeService(BaseService):

    def __init__(self):
        self.log = self.create_logger('knowledge_svc')
        self.schema = dict(facts=[], relationships=[], rules=[], constraints=dict())
        self.fact_ram = copy.deepcopy(self.schema)

    # -- Fact API --
    async def _add_fact(self, fact, constraints=None):
        """
        Add a fact to the internal store
        :param fact: Fact to add
        :param constraints: any potential constraints
        """
        # facts can now be highly controlled, with visibility at the operation level, agent(s) level, or
        # custom groupings
        if not any(x == fact for x in self.fact_ram['facts']):
            fact._knowledge_id = uuid.uuid4()
            self.fact_ram['facts'].append(fact)
            if constraints:
                self.fact_ram['constraints'][fact._knowledge_id] = constraints

    async def _update_fact(self, criteria, updates):
        """
        Update a fact in the internal store
        :param criteria: dictionary containing fields to match on
        :param updates: dictionary containing fields to replace
        """
        matches = await self._get_facts(criteria)
        for match in matches:
            for k, v in updates.items():
                if getattr(match, k, False):
                    setattr(match, k, v)

    async def _get_facts(self, criteria, restrictions=None):
        """
        Retrieve a fact from the internal store
        :param criteria: dictionary containing fields to match on
        :return: list of facts matching the criteria
        """
        complete_list = await self._locate('facts', query=criteria)
        return await self._apply_restrictions(complete_list, restrictions)

    async def _delete_fact(self, criteria):
        """
        Delete a fact from the internal store
        :param criteria: dictionary containing fields to match on
        """
        return await self._remove('facts', criteria)

    async def _get_meta_facts(self, meta_fact=None, agent=None, group=None):
        # Returns the complete set of facts associated with a meta-fact construct
        raise NotImplementedError

    async def _get_fact_origin(self, fact):
        """
        Identify the place where a fact originated, either the source that loaded it or its original link
        :param fact: Fact to get origin for (can be either a trait string or a full blown fact)
        :return: tuple - (String of either origin source id or origin link id, fact origin type)
        """
        workspace = copy.deepcopy(fact)
        if not getattr(workspace, 'links', False):
            fact_search = await self._get_facts(dict(trait=workspace))
            if fact_search:
                workspace = fact_search[0]
            else:
                return None, None

        if workspace.links:
            return str(workspace.links[0]), workspace.origin_type  # Return the id of the first link associated
        elif workspace.source:
            return str(workspace.source), workspace.origin_type  # Return the source of the fact if not found

        return None, None  # Default return value

    # -- Relationships API --

    async def _get_relationships(self, criteria, restrictions=None):
        """
        Retrieve relationships from the internal store
        :param criteria: dictionary containing fields to match on
        :return: list of matching relationships
        """
        complete_list = await self._locate('relationships', query=criteria)
        return await self._apply_restrictions(complete_list, restrictions)

    async def _add_relationship(self, relationship, constraints=None):
        """
        Add a relationship to the internal store
        :param relationship: Relationship object to add
        :param constraints: optional constraints on the use of the relationship
        """
        if not any((x.source == relationship.source) and (x.edge == relationship.edge) and
                   (x.target == relationship.target) and (x.origin == relationship.origin or not relationship.origin)
                   for x in self.fact_ram['relationships']):
            relationship._knowledge_id = uuid.uuid4()
            self.fact_ram['relationships'].append(relationship)
            if constraints:
                self.fact_ram['constraints'][relationship._knowledge_id] = constraints

    async def _update_relationship(self, criteria, updates):
        """
        Update a relationship in the internal store
        :param criteria: dictionary containing fields to match on
        :param updates: dictionary containing fields to modify
        """
        matches = await self._get_relationships(criteria)
        for match in matches:
            for k, v in updates.items():
                if getattr(match, k, "eMpTy") != "eMpTy":
                    if isinstance(getattr(match, k), Fact) and not isinstance(v, Fact):
                        handle = getattr(match, k)
                        for x in v.keys():
                            setattr(handle, x, v[x])
                    else:
                        setattr(match, k, v)

    async def _delete_relationship(self, criteria):
        """
        Remove a relationship from the internal store
        :param criteria: dictionary containing fields to match on
        """
        return await self._remove('relationships', criteria)

    # --- Rule API ---

    async def _add_rule(self, rule, constraints=None):
        """
        Add a rule to the internal store
        :param rule: Rule object to add
        :param constraints: dictionary containing fields to match on
        """
        ###
        #    rule.action: [DENY, ALLOW, EXCLUSIVE, EXCLUSIVE_TRAIT, EXCLUSIVE_VALUE], 'EXCLUSIVE_*' actions denote that
        #        the trait/value will be made mutually exclusive in its use to the agent/group/operation that is
        #        specified for. Essentially a fact is binded to mutex, and only one action can be using the fact
        #        at any one time.
        ###
        if not any((x.action == rule.action) and (x.trait == rule.trait) for x in self.fact_ram['rules']):
            rule._knowledge_id = uuid.uuid4()
            self.fact_ram['rules'].append(rule)
            if constraints:
                self.fact_ram['constraints'][rule._knowledge_id] = constraints

    async def _get_rules(self, criteria, restrictions=None):
        """
        Retrieve rules from the internal store
        :param criteria: dictionary containing fields to match on
        :return: list of matching rules
        """
        complete_list = [x for x in self.fact_ram['rules'] if await self._check_rule(x, criteria, True)]
        return await self._apply_restrictions(complete_list, restrictions)

    async def _delete_rule(self, criteria):
        """
        Remove a rule from the internal store
        :param criteria: dictionary containing fields to match on
        """
        sublist = [x for x in self.fact_ram['rules'] if await self._check_rule(x, criteria)]
        await self._clear_matching_constraints(sublist)
        self.fact_ram['rules'][:] = [x for x in self.fact_ram['rules'] if x not in sublist]
        return

    @staticmethod
    async def _check_rule(rule, desired_quals, wildcard=False):
        """
        Check whether or not a rule matches the provided criteria
        :param rule: Rule object to check
        :param desired_quals: dictionary containing fields to match on
        :param wildcard: whether or not to do wildcard matching on the 'match' field
        :return:
        """
        criteria_matches = []
        for k, v in desired_quals.items():
            if getattr(rule, k) == v:
                criteria_matches.append(True)
            elif k == 'match' and wildcard:  # support wildcard matching
                val = getattr(rule, k)
                local_stub = val.split('*')
                input_stub = v.split('*')
                if all(x in v for x in local_stub) or all(x in val for x in input_stub):
                    criteria_matches.append(True)
        if len(criteria_matches) == len(desired_quals):
            return True

    async def _locate(self, object_name, query=None):
        """
        Locate a matching object in the internal store
        :param object_name: object type
        :param query: dictionary of fields to match on
        :return: list of matching objects
        """
        try:
            return [obj for obj in self.fact_ram[object_name] if self._wildcard_match(obj, query)]
        except Exception as e:
            self.log.error('[!] LOCATE: %s' % e)

    async def _remove(self, object_name, query):
        """
        Remove objects from the internal store
        :param object_name: object type
        :param query: dictionary of fields to match on
        """
        try:
            sublist = [obj for obj in self.fact_ram[object_name] if self._wildcard_match(obj, query)]
            await self._clear_matching_constraints(sublist)
            self.fact_ram[object_name][:] = [obj for obj in self.fact_ram[object_name] if obj not in sublist]
        except Exception as e:
            self.log.error('[!] REMOVE: %s' % e)

    async def _clear_matching_constraints(self, objs):
        """
        Remove constraints associated with objects as part of deletion
        :param objs: list of objects being removed
        """
        for obj in objs:
            if obj._knowledge_id in self.fact_ram['constraints']:
                del self.fact_ram['constraints'][obj._knowledge_id]

    async def _get_matching_constraints(self, objs):
        """
        Identify matching constraints associated with objects
        :param objs: list of objects to get constraints for
        :return: list of relevant constraints
        """
        k_ids = [obj._knowledge_id for obj in objs]
        return [self.fact_ram['constraints'][constraint] for constraint in self.fact_ram['constraints']
                if constraint in k_ids]

    @staticmethod
    def _delete_file(path):
        if not os.path.exists(path):
            return
        elif os.path.isdir(path):
            shutil.rmtree(path)
        else:
            os.remove(path)

    @staticmethod
    def _destroy():
        """
        Reset the caldera data directory and server state.

        This creates a gzipped tarball backup of the data files tracked by caldera.
        Paths are preserved within the tarball, with all files having "data/" as the
        root.

        :return: None
        """
        if not os.path.exists(DATA_BACKUP_DIR):
            os.mkdir(DATA_BACKUP_DIR)

        timestamp = datetime.utcnow().strftime('%Y%m%d%H%M%S')
        tarball_path = os.path.join(DATA_BACKUP_DIR, f'backup-{timestamp}.tar.gz')

        with tarfile.open(tarball_path, 'w:gz') as tarball:
            if os.path.isfile(FACT_STORE_PATH):
                tarball.add(FACT_STORE_PATH)
                BaseKnowledgeService._delete_file(FACT_STORE_PATH)

    async def _save_state(self):
        """
        Saves the current internal store state to disk

        :return: None
        """
        await self.get_service('file_svc').save_file(FACT_STORE_PATH.split(os.path.sep)[1], pickle.dumps(self.fact_ram),
                                                     FACT_STORE_PATH.split(os.path.sep)[0])

    async def _restore_state(self):
        """
        Loads the internal store state from disk

        :return: None
        """
        if os.path.exists(FACT_STORE_PATH):
            _, store = await self.get_service('file_svc').read_file(FACT_STORE_PATH.split(os.path.sep)[1],
                                                                    FACT_STORE_PATH.split(os.path.sep)[0])
            # Pickle is only used to load a local file that caldera creates. Pickled data is not
            # received over the network.
            ram = pickle.loads(store)  # nosec
            for key in ram.keys():
                self.fact_ram[key] = []
                for c_object in ram[key]:
                    handle = self._load_wrapper(key)
                    constraints = []
                    if c_object._knowledge_id in ram[key]:
                        constraints = ram[key][c_object._knowledge_id]
                    await handle(c_object, constraints=constraints)
            self.log.debug('Restored data from persistent storage')

    def _load_wrapper(self, key):
        """
        Support wrapper to process different object types during load
        :param key: type of object to load
        :return: function handle to the correct loader function for the associated object type
        """
        if key == 'facts':
            return self._add_fact
        elif key == 'relationships':
            return self._add_relationship
        elif key == 'rules':
            return self._add_rule

    async def _apply_restrictions(self, complete_list, restrictions):
        """
        Apply restrictions to a list of objects (type agnostic)
        :param complete_list: List of objects to apply restrictions to
        :param restrictions: The restrictions to be applied (checked against the constraints internal store)
        :return: filtered list of input objects
        """
        def _check_restrictions(existing_limitations, desired_limitations):
            for restriction_name, restriction_value in desired_limitations:
                if restriction_name in existing_limitations:
                    if existing_limitations[restriction_name] != restriction_value:
                        return False
            return True

        if not restrictions:
            return complete_list
        ret = []
        for entry in complete_list:
            constraints = await self._get_matching_constraints([entry])
            if _check_restrictions(constraints, restrictions):
                ret.append(entry)
        return ret

    def _wildcard_match(self, obj, criteria):
        """
        Modified variant of the normal `match` method for objects that will return matches for wildcard fields
            [fact object].source
            [relationship object].origin
        :param obj: The object to validate
        :param criteria: The values to check against
        """
        if not criteria:
            return obj
        criteria_matches = []
        for k, v in criteria.items():
            if type(v) is tuple:
                for val in v:
                    if getattr(obj, k) == val:
                        criteria_matches.append(True)
            else:
                if getattr(obj, k) == v:
                    criteria_matches.append(True)
                elif isinstance(getattr(obj, k), Fact):  # this is a match based on a fact object in a relationship
                    if isinstance(v, Fact):
                        if getattr(obj, k) == v:
                            criteria_matches.append(True)
                    elif self._wildcard_match(getattr(obj, k), v):
                        criteria_matches.append(True)
                else:  # Wildcard match check
                    if (k == 'source' and isinstance(obj, Fact)) or (k == 'origin' and isinstance(obj, Relationship)):
                        if getattr(obj, k) == WILDCARD_STRING:
                            criteria_matches.append(True)
        if len(criteria_matches) >= len(criteria) and all(criteria_matches):
            return obj
