import sys
from typing import (
    Optional,
)

try:
    import statsd
except ImportError:
    statsd = None

from galaxy.util import asbool


# TODO: optimize with two separate implementations around statsd_influxdb?
class VanillaGalaxyStatsdClient:
    def __init__(self, statsd_host, statsd_port, statsd_prefix, statsd_influxdb, statsd_mock_calls=False):
        if not statsd:
            raise ImportError(
                "Statsd logging configured, but no statsd python module found. "
                "Please install the python statsd module to use this functionality."
            )

        self.metric_infix = ""
        self.statsd_influxdb = asbool(statsd_influxdb)
        if self.statsd_influxdb:
            statsd_prefix = statsd_prefix.strip(",")
        if statsd_mock_calls:
            self.statsd_client = MockStatsClient()
        else:
            self.statsd_client = statsd.StatsClient(statsd_host, statsd_port, prefix=statsd_prefix)

    def timing(self, path, time, tags=None):
        infix = self._effective_infix(path, tags)
        self.statsd_client.timing(infix + path, time)

    def incr(self, path, n=1, tags=None):
        infix = self._effective_infix(path, tags)
        self.statsd_client.incr(infix + path, n)

    def _effective_infix(self, path, tags):
        tags = tags or {}
        if self.statsd_influxdb and tags:
            return f",{','.join(f'{k}={v}' for k, v in tags.items())}" + ",path="
        if self.statsd_influxdb:
            return ",path="
        else:
            return ""


CURRENT_TEST: Optional[str] = None
CURRENT_TEST_METRICS: Optional[dict[str, dict]] = None


class PyTestGalaxyStatsdClient(VanillaGalaxyStatsdClient):
    def timing(self, path, time, tags=None):
        if (metrics := CURRENT_TEST_METRICS) is not None:
            timing = metrics["timing"]
            if path not in timing:
                timing[path] = []
            timing[path].append({"time": time, "tags": tags})
        super().timing(path, time, tags=tags)

    def incr(self, path, n=1, tags=None):
        if (metrics := CURRENT_TEST_METRICS) is not None:
            counter = metrics["counter"]
            if path not in counter:
                counter[path] = []
            counter[path].append({"n": n, "tags": tags})
        super().incr(path, n=n, tags=tags)

    def _effective_infix(self, path, tags):
        if (current_test := CURRENT_TEST) is not None:
            tags = tags or {}
            tags["test"] = current_test
        return super()._effective_infix(path, tags)


class MockStatsClient:
    def timing(self, path, time, tags=None):
        pass

    def incr(self, path, n=1, tags=None):
        pass


# Replace stats collector if in pytest environment
if "pytest" in sys.modules:
    GalaxyStatsdClient = PyTestGalaxyStatsdClient
else:
    GalaxyStatsdClient = VanillaGalaxyStatsdClient  # type: ignore[assignment,misc]


__all__ = ("GalaxyStatsdClient",)
