from typing import Union,List,Dict,Any,Callable,Tuple,Optional

# External Modules
import traceback
import networkx as nx         # type: ignore
from copy   import deepcopy
from pprint import pformat
# Iternal Modules
from dbgen.support.datatypes.func     import Func
from dbgen.support.datatypes.misc     import ExternalError,Arg
from dbgen.support.utils              import (mkInsCmd,sqlite_execute
                                                    ,func_to_ind,sqlite_select
                                                    ,topsort_with_dict)

##############################################################################
class Temp(object):
    """
    - A Template (Temp) object specifies a composition of functions which
        expand the namespace generated by a query.

    - These are related by a set of function *instances* (a function can be
        instantiated multiple times with different arguments).

    - For every function instance, we need to specify where it gets its inputs
        by referring to the namespace at that moment in time.
    """

    def __init__(self
                ,tempfuncs : List['TempFunc']   = []
                ,constants : Dict[str,'Const'] = {}
                ) -> None:
        self.tempfuncs = self.order_tempfuncs(tempfuncs)
        self.constants = constants

    def src(self)->str:
        tfs   = '\n\t\t,'.join([tf.src() for tf in self.tempfuncs])
        cnsts = '\n\t\t,'.join(["'%s':%s"%(k,v.src()) for k,v in self.constants.items()])
        return "Temp(constants = {%s}\n\t\t,tempfuncs=[%s])"%(cnsts,tfs)

    def __str__(self)->str:
        return pformat(self.__dict__)

    def __repr__(self)->str:
        return str(self)

    def __eq__(self,other : object)->bool:
        return self.__dict__==other.__dict__

    @staticmethod
    def order_tempfuncs(tfs:List['TempFunc'])->List['TempFunc']:
        """
        Prior to storing the list of tempfuncs as a field of the Template,
        order them such that any TF that calls some other TF occurs afterwards
        """

        G = nx.DiGraph() # each template has a DAG, with tempfuncs as nodes
        G.add_nodes_from([tf.name for tf in tfs])

        tf_dict = {tf.name : tf for tf in tfs}

        for tf in tfs:
            for tfa in tf.args:
                if tfa.name in tf_dict.keys():
                    G.add_edge(tfa.name,tf.name)
        return topsort_with_dict(G,tf_dict)

    def apply_template(self
                      ,query_dic  : Dict[str,Any]
                      ) -> Dict[str,Any]:
        """
        This should really be done by creating a computational graph...(DART?)
        """
        consts = {k:v.val for k,v in self.constants.items()}
        dic = {**consts,**query_dic} # MERGE

        for tf in self.tempfuncs:
            dic[tf.name] = tf.apply_func(dic) #type: ignore

        return dic

    def add_temp(self
                ,sqlite_pth   : str
                ,rule_id : int
                ) -> None:
        """
        Add to sqlite meta.db
        """
        for i,tf in enumerate(self.tempfuncs):
            tf.add_tempfunc(sqlite_pth,rule_id,i)
        for j,(constname,const) in enumerate(self.constants.items()):
            const.add_const(sqlite_pth,rule_id,j,constname)

#############################################################################
class TempFunc(object):
    """
    - This object represents a particular function call.
    - Should be given a unique name
    - Each arg refers to the namespace (option of indexing a tuple)
    """
    def __init__(self
                ,func     : Callable
                ,name     : Optional[str] = None
                ,args     : List[Arg]     = []
                ) -> None:
        assert callable(func), 'Not a func: '+str(func)
        self.func     = Func(func)

        if name is None:
            name = self.func.name

        self.name     = name
        self.args     = args

    def src(self)->str:
        args = '\n\t\t,'.join(a.src() for a in self.args)
        fmtargs = [self.name,self.name,args]
        return "TempFunc(funcname='{0}'\n\t\t\t,name='{1}'\n\t\t\t,args=[{2}])".format(*fmtargs)

    def __str__(self)->str:
        return pformat(self.__dict__)

    def __repr__(self)->str:
        return str(self)

    def __eq__(self,other : object)->bool:
        return self.__dict__==other.__dict__

    def apply_func(self
                  ,curr_dict : Dict[str,Any]
                  ) -> Dict[str,Any]:
        """
        Take a TempFunc's function and wrap it such that it can accept a namespace
            dictionary. The information about how the function interacts with the
            namespace is contained in the TempFunc's args.
        """
        inputvars = [arg.arg_get(curr_dict) for arg in self.args]

        try:
            return self.func.apply(*inputvars)
        except Exception as e:
            msg = '\tApplying func %s in tempfunc %s:\n\t'%(self.func.name,self.name)
            raise ExternalError(msg + traceback.format_exc())

    def add_tempfunc(self
                    ,sqlite_pth  : str
                    ,trans_id    : int
                    ,tf_id       : int
                    ) -> None:
        """
        Add to meta.db
        """
        func_id = func_to_ind(sqlite_pth,self.name) # type: ignore
        cols    = ['rule_id','tempfunc_id','func_id','name']
        q       = mkInsCmd('tempfunc',cols,sqlite=True)
        sqlite_execute(sqlite_pth,q,[trans_id,tf_id,func_id,self.name])
        for i,a in enumerate(self.args):
            a.add_arg(sqlite_pth,'temp',trans_id,tf_id,i)

    @staticmethod
    def maybe_index(d   : dict
                   ,arg : Tuple[str,Optional[int]]
                   ) -> Any:
        """
        If TempFunc arg specifies an index, then we are expecting a tuple and index it
        """
        if arg[1] is None:
            return d[arg[0]]
        else:
            return d[arg[0]][arg[1]]

class Const(object):
    """
    Embed a constant value of any type into the namespace for a given rule.
    """
    def __init__(self,val:Any)->None:
        self.val = val

    def src(self)->str:
        return 'Const(%s)'%repr(self.val)

    def __str__(self)->str:
        return pformat(self.__dict__)

    def __repr__(self)->str:
        return str(self)

    def __eq__(self,other : object)->bool:
        return self.__dict__==other.__dict__

    def get_const(self,*args: Any)->Any:
        return self.val

    def get_datatype(self)->str:
        return str(self.val.__class__)

    def add_const(self
                 ,sqlite_pth : str
                 ,trans_id   : int
                 ,const_id   : int
                 ,const_name : str
                 )->None:
        """
        Add to meta.db
        """
        datatype = self.get_datatype()
        cols     = ['rule_id','const_id','name','datatype','val']
        q        = mkInsCmd('constants',cols,sqlite=True)
        sqlite_execute(sqlite_pth,q,[trans_id,const_id,const_name,datatype,str(self.val)])

################################################################################################
################################################################################################

#########################
# Special Temp shortcuts
#-----------------------
def identity(x:Any)->Any:
    return x

def SimpleTemp(func    : Callable
              ,inputs  : List[str] = []
              ,outputs : List[str] = []
              ,const   : Dict[str,Const] = {}
              ) -> Temp:
    """
    For a rule that has one function in its template
    """
    main_func = TempFunc(func,'main',args=[Arg(x) for x in inputs])
    if len(outputs)==1:
        exported = [TempFunc(identity,outputs[0],[Arg('main')])] # don't index if output is not expected to be a tuple
    else:
        exported  = [TempFunc(identity,x,[Arg('main',i)])
                        for i,x in enumerate(outputs)]
    return Temp(tempfuncs = [main_func]+exported
               ,constants = const)

def SimplePipe(funcs   : List[Callable]
              ,inputs  : List[str]
              ,outputs : List[str]
              ,const   : Dict[str,Const] = {}
              ) -> Temp:
    """
    A series of functions where the output of each function is fed to the
    following function in the pipeline. You must specify the names of the
    inputs and outputs to the whole process
    """
    init_func = TempFunc(funcs[0],args=[Arg(x) for x in inputs])
    remaining = [TempFunc(funcs[i],args=[Arg(funcs[i-1].__name__)]) for i in range(1,len(funcs))]
    if len(outputs)==1:
        exported = [TempFunc(identity,outputs[0],[Arg(funcs[-1].__name__)])] # don't index if output is not expected to be a tuple
    else:
        exported  = [TempFunc(identity,x,[Arg(funcs[-1].__name__,i)])
                        for i,x in enumerate(outputs)]
    return Temp(tempfuncs = [init_func] + remaining + exported
               ,constants = const)
