import iso8601
from typing import Optional, Callable, Dict, Any, List
from inspect import signature
from functools import wraps
from datetime import timedelta

from aw_core.models import Event
from aw_datastore import Datastore

from aw_transform import (
    filter_period_intersect,
    filter_keyvals,
    filter_keyvals_regex,
    period_union,
    union_no_overlap,
    categorize,
    tag,
    Rule,
    merge_events_by_keys,
    chunk_events_by_key,
    sort_by_timestamp,
    sort_by_duration,
    sum_durations,
    concat,
    split_url_events,
    simplify_string,
    flood,
    limit_events,
)

from .exceptions import QueryFunctionException


def _verify_bucket_exists(datastore, bucketname):
    if bucketname in datastore.buckets():
        return
    else:
        raise QueryFunctionException(f"There's no bucket named '{bucketname}'")


def _verify_variable_is_type(variable, t):
    if not isinstance(variable, t):
        raise QueryFunctionException(
            "Variable '{}' passed to function call is of invalid type. Expected {} but was {}".format(
                variable, t, type(variable)
            )
        )


# TODO: proper type checking (typecheck-decorator in pypi?)


TNamespace = Dict[str, Any]
TQueryFunction = Callable[..., Any]


"""
    Declarations
"""
functions: Dict[str, TQueryFunction] = {}


def q2_function(transform_func=None):
    """
    Decorator used to register query functions.

    Automatically adds mock arguments for Datastore and TNamespace
    if not in function signature.
    """

    def h(f):
        sig = signature(f)
        # If function lacks docstring, use docstring from underlying function in aw_transform
        if transform_func and transform_func.__doc__ and not f.__doc__:
            f.__doc__ = ".. note:: Documentation automatically copied from underlying function `aw_transform.{func_name}`\n\n{func_doc}".format(
                func_name=transform_func.__name__, func_doc=transform_func.__doc__
            )

        @wraps(f)
        def g(datastore: Datastore, namespace: TNamespace, *args, **kwargs):
            # Remove datastore and namespace argument for functions that don't need it
            args = (datastore, namespace, *args)
            if TNamespace not in (sig.parameters[p].annotation for p in sig.parameters):
                args = (args[0], *args[2:])
            if Datastore not in (sig.parameters[p].annotation for p in sig.parameters):
                args = args[1:]
            return f(*args, **kwargs)

        fname = f.__name__
        if fname[:3] == "q2_":
            fname = fname[3:]
        functions[fname] = g
        return g

    return h


def q2_typecheck(f):
    """Decorator that typechecks using `_verify_variable_is_type`"""
    sig = signature(f)

    @wraps(f)
    def g(*args, **kwargs):
        # FIXME: If the first argument passed to a query2 function is a straight [] then the second argument disappears from the argument list for unknown reasons, which breaks things
        for i, p in enumerate(sig.parameters):
            param = sig.parameters[p]

            # print(f"Checking that param ({param}) was {param.annotation}, value: {args[i]}")
            # FIXME: Won't check keyword arguments
            if (
                param.annotation in [list, str, int, float]
                and param.default == param.empty
            ):
                _verify_variable_is_type(args[i], param.annotation)

        return f(*args, **kwargs)

    return g


"""
    Getting buckets
"""


@q2_function()
@q2_typecheck
def q2_find_bucket(
    datastore: Datastore, filter_str: str, hostname: Optional[str] = None
):
    """Find bucket by using a filter_str (to avoid hardcoding bucket names)"""
    for bucket in datastore.buckets():
        if filter_str in bucket:
            bucket_metadata = datastore[bucket].metadata()
            if hostname:
                if bucket_metadata["hostname"] == hostname:
                    return bucket
            else:
                return bucket
    raise QueryFunctionException(
        "Unable to find bucket matching '{}' (hostname filter set to '{}')".format(
            filter_str, hostname
        )
    )


"""
    Data gathering functions
"""


@q2_function()
@q2_typecheck
def q2_query_bucket(
    datastore: Datastore, namespace: TNamespace, bucketname: str
) -> List[Event]:
    _verify_bucket_exists(datastore, bucketname)
    try:
        starttime = iso8601.parse_date(namespace["STARTTIME"])
        endtime = iso8601.parse_date(namespace["ENDTIME"])
    except iso8601.ParseError:
        raise QueryFunctionException(
            "Unable to parse starttime/endtime for query_bucket"
        )
    return datastore[bucketname].get(starttime=starttime, endtime=endtime)


@q2_function()
@q2_typecheck
def q2_query_bucket_eventcount(
    datastore: Datastore, namespace: TNamespace, bucketname: str
) -> int:
    _verify_bucket_exists(datastore, bucketname)
    starttime = iso8601.parse_date(namespace["STARTTIME"])
    endtime = iso8601.parse_date(namespace["ENDTIME"])
    return datastore[bucketname].get_eventcount(starttime=starttime, endtime=endtime)


"""
    Filtering functions
"""


@q2_function(filter_keyvals)
@q2_typecheck
def q2_filter_keyvals(events: list, key: str, vals: list) -> List[Event]:
    return filter_keyvals(events, key, vals, False)


@q2_function(filter_keyvals)
@q2_typecheck
def q2_exclude_keyvals(events: list, key: str, vals: list) -> List[Event]:
    return filter_keyvals(events, key, vals, True)


@q2_function(filter_keyvals_regex)
@q2_typecheck
def q2_filter_keyvals_regex(events: list, key: str, regex: str) -> List[Event]:
    return filter_keyvals_regex(events, key, regex)


@q2_function(filter_period_intersect)
@q2_typecheck
def q2_filter_period_intersect(events: list, filterevents: list) -> List[Event]:
    return filter_period_intersect(events, filterevents)


@q2_function(period_union)
@q2_typecheck
def q2_period_union(events1: list, events2: list) -> List[Event]:
    return period_union(events1, events2)


@q2_function(limit_events)
@q2_typecheck
def q2_limit_events(events: list, count: int) -> List[Event]:
    return limit_events(events, count)


"""
    Merge functions
"""


@q2_function(merge_events_by_keys)
@q2_typecheck
def q2_merge_events_by_keys(events: list, keys: list) -> List[Event]:
    return merge_events_by_keys(events, keys)


@q2_function(chunk_events_by_key)
@q2_typecheck
def q2_chunk_events_by_key(events: list, key: str) -> List[Event]:
    return chunk_events_by_key(events, key)


"""
    Sort functions
"""


@q2_function(sort_by_timestamp)
@q2_typecheck
def q2_sort_by_timestamp(events: list) -> List[Event]:
    return sort_by_timestamp(events)


@q2_function(sort_by_duration)
@q2_typecheck
def q2_sort_by_duration(events: list) -> List[Event]:
    return sort_by_duration(events)


"""
    Summarizing functions
"""


@q2_function(sum_durations)
@q2_typecheck
def q2_sum_durations(events: list) -> timedelta:
    return sum_durations(events)


@q2_function(concat)
@q2_typecheck
def q2_concat(events1: list, events2: list) -> List[Event]:
    return concat(events1, events2)


@q2_function(union_no_overlap)
@q2_typecheck
def q2_union_no_overlap(events1: list, events2: list) -> List[Event]:
    return union_no_overlap(events1, events2)


"""
    Flood functions
"""


@q2_function(flood)
@q2_typecheck
def q2_flood(events: list) -> List[Event]:
    return flood(events)


"""
    Watcher specific functions
"""


@q2_function(split_url_events)
@q2_typecheck
def q2_split_url_events(events: list) -> List[Event]:
    return split_url_events(events)


@q2_function(simplify_string)
@q2_typecheck
def q2_simplify_window_titles(events: list, key: str) -> List[Event]:
    return simplify_string(events, key=key)


"""
    Test functions
"""


@q2_function()
@q2_typecheck
def q2_nop():
    """No operation function for unittesting"""
    return 1


"""
    Classify
"""


@q2_function(categorize)
@q2_typecheck
def q2_categorize(events: list, classes: list):
    classes = [(_cls, Rule(rule_dict)) for _cls, rule_dict in classes]
    return categorize(events, classes)


@q2_function(tag)
@q2_typecheck
def q2_tag(events: list, classes: list):
    classes = [(_cls, Rule(rule_dict)) for _cls, rule_dict in classes]
    return tag(events, classes)
