
import aiohttp.client as aiohttp
import aiohttp.tracing as tracing
from backup.config import Config, Setting
from .logger import getLogger

logger = getLogger(__name__)


class TracingSession(aiohttp.ClientSession):
    def __init__(self, config: Config, **kwargs):
        self.config: Config = config
        self.trace_config = aiohttp.TraceConfig()
        if self.config.get(Setting.TRACE_REQUESTS):
            self.trace_config.on_request_start.append(self.trace_request_start)
            self.trace_config.on_request_end.append(self.trace_request_end)
            self.trace_config.on_request_exception.append(self.trace_request_exception)
            self.trace_config.on_response_chunk_received.append(self.trace_chunk_recv)
            self.trace_config.on_request_chunk_sent.append(self.trace_chunk_sent)
        self._record = False
        self._records = []
        super().__init__(**kwargs, trace_configs=[self.trace_config])

    @property
    def record(self) -> bool:
        return self._record

    @record.setter
    def record(self, value: bool):
        self._record = value

    def clearRecord(self):
        self._records = []

    async def _request(self, method, url, **kwargs):
        if self.record:
            self._records.append({'method': method, 'url': url, 'other_args': kwargs})
        if not self.config.get(Setting.TRACE_REQUESTS):
            return await super()._request(method, url, **kwargs, trace_request_ctx="{0} {1}".format(method, url))
        resp = await super()._request(method, url, **kwargs, trace_request_ctx="{0} {1}".format(method, url))
        logger.trace("Initial response data: %s %s", resp.method, resp.url)
        logger.trace("  headers:")
        for header in resp.headers:
            logger.trace("    %s: %s", header, resp.headers[header])
        return resp

    def print_headers(self, headers):
        for header in headers:
            logger.trace("    %s: %s", header, 'redacted' if ('Authorization' in header) else headers[header])

    async def trace_request_start(self, session, trace_config_ctx, params: tracing.TraceRequestStartParams):
        logger.trace("Request started: %s", trace_config_ctx.trace_request_ctx)
        logger.trace("  sent headers:")
        self.print_headers(params.headers)

    async def trace_request_exception(self, session, trace_config_ctx, params: tracing.TraceRequestExceptionParams):
        logger.trace("Request exception: %s", trace_config_ctx.trace_request_ctx)
        logger.trace(logger.formatException(params.exception))

    async def trace_request_end(self, session, trace_config_ctx, params: tracing.TraceRequestEndParams):
        logger.trace("Request ending: %s", trace_config_ctx.trace_request_ctx)
        logger.trace("  status: %s", params.response.status)
        logger.trace("  recieved headers:")
        self.print_headers(params.response.headers)

    async def trace_chunk_recv(self, session, trace_config_ctx, params: tracing.TraceResponseChunkReceivedParams):
        logger.trace("Recieved chunk: %s", trace_config_ctx.trace_request_ctx)
        logger.trace("  length: %s", len(params.chunk))

    async def trace_chunk_sent(self, session, trace_config_ctx, params: tracing.TraceRequestChunkSentParams):
        logger.trace("Sent chunk: %s", trace_config_ctx.trace_request_ctx)
        logger.trace("  length: %s", len(params.chunk))
