import time

from requests_futures.sessions import FuturesSession


class ElapsedFuturesSession(FuturesSession):
    def request(self, method, url, hooks=None, *args, **kwargs):
        start = time.time()
        if hooks is None:
            hooks = {}

        def timing(r, *args, **kwargs):
            r.elapsed = round((time.time() - start) * 1000, 2)

        try:
            if isinstance(hooks["response"], (list, tuple)):
                # needs to be first so we don't time other hooks execution
                hooks["response"].insert(0, timing)
            else:
                hooks["response"] = [timing, hooks["response"]]
        except KeyError:
            hooks["response"] = timing

        return super(ElapsedFuturesSession, self).request(
            method, url, hooks=hooks, *args, **kwargs
        )