from functools import wraps

import asyncio
try:
    from aiomeasures import StatsD
except ImportError:
    StatsD = None
from alamo_common.conf import AlamoSettings
from alamo_common import stats


settings = AlamoSettings()
_aiostatsd = None


def setup():
    global _aiostatsd
    host = 'udp://{}:{}'.format(settings.STATSD_HOST,
                                settings.STATSD_PORT)
    try:
        _aiostatsd = StatsD(
            addr=host,
            prefix=settings.STATSD_PREFIX
        )
    except TypeError:
        raise RuntimeError('Could not use StatsD from aiomeasures.')


def get_aiostatsd():
    if _aiostatsd is None:
        setup()
    return _aiostatsd


class increment(stats.StatMixin):
    """Statsd decorator and context manager to `incr` stat.

    >>> from alamo_common.aiostats import increment
    ...
    ... @increment('stat_name')
    ... def incremented():
    ...   pass
    ...
    ... with increment(metric_name='stat'):
    ...   pass
    ...
    ... increment.incr('stat', value=2)
    """

    def __call__(self, f):
        if asyncio.iscoroutinefunction(f):
            @wraps(f)
            async def wrapped(*args, **kwargs):
                get_aiostatsd().incr(self.get_metric(f.__name__), 1)
                return await f(*args, **kwargs)
        else:
            @wraps(f)
            def wrapped(*args, **kwargs):
                get_aiostatsd().incr(self.get_metric(f.__name__), 1)
                return f(*args, **kwargs)

        return wrapped

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        get_aiostatsd().incr(self.metric_name or 'unknown', 1)

    @classmethod
    def incr(cls, stat, value=1, hostname=None):
        """Increment a stat by `count`.

        :param stat: stat name
        :param value: value
        :param hostname: if hostname it send stat with hostname prefix
        """
        klass = cls(metric_name=stat, hostname=hostname)
        get_aiostatsd().incr(klass.get_metric(''), value=value)

    @classmethod
    def decr(cls, stat, value=1, hostname=None):
        """Decrement a stat by `count`.

        :param stat: stat name
        :param value: value
        :param hostname: if hostname it send stat with hostname prefix
        """
        klass = cls(metric_name=stat, hostname=hostname)
        get_aiostatsd().decr(klass.get_metric(''), value=value)


class timer(stats.StatMixin):
    """Statsd  decorator and context manager to timed method execution.

    >>> from alamo_common.aiostats import timer
    ...
    ... @timer(hostname=True)
    ... def timed_method():
    ...   # do something
    ...
    ... with timer(metric_name='timed_stat'):
    ...   # do something
    ...
    ... timing = timer.start()
    ... # do something
    ... timing.stop()
    """
    timer = None

    def __call__(self, f):
        if asyncio.iscoroutinefunction(f):
            @wraps(f)
            async def wrapped(*args, **kwargs):
                with get_aiostatsd().timer(self.get_metric(f.__name__)):
                    return await f(*args, **kwargs)
        else:
            @wraps(f)
            def wrapped(*args, **kwargs):
                with get_aiostatsd().timer(self.get_metric(f.__name__)):
                    return f(*args, **kwargs)

        return wrapped

    def __enter__(self):
        self.timer = get_aiostatsd().timer(self.metric_name)
        self.timer.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.timer.stop()

    @classmethod
    def start(cls, metric_name, hostname=None):
        """Manually start timer."""
        klass = cls(metric_name=metric_name, hostname=hostname)
        klass.timer = get_aiostatsd().timer(klass.get_metric(''))
        return klass

    def stop(self):
        """Stop and send stats to statsd."""
        if self.timer is None:
            raise RuntimeError('Timer has not been created.')
        self.timer.stop()
