import logging
import textwrap
from collections import defaultdict
from contextlib import ContextDecorator, ExitStack

from django.conf import settings

from .capture import native_query_capture

logger = logging.getLogger(__name__)


class QueryCountExceededException(Exception):
    def __init__(
        self,
        *,
        expected_query_count,
        actual_query_count,
        extra_queries,
        affected_cls_or_fn_name=None,
        view_cls_name=None,
        request=None,
    ):
        self.expected_query_count = expected_query_count
        self.actual_query_count = actual_query_count
        self.extra_queries = extra_queries
        self.extra_queries_pretty_str = self._prettify_extra_queries(self.extra_queries)
        message_lines = [
            f"Possible N+1 problem {('on ' + affected_cls_or_fn_name) or ''}",
            (f"View={view_cls_name}" if view_cls_name else ""),
            (f"Endpoint={request.path}" if request else ""),
            (f"Method={request.method}" if request else ""),
            (
                f"Exceeded expected queries, expected={self.expected_query_count}, "
                f"actual={self.actual_query_count}, "
                f"queries={self.extra_queries_pretty_str}"
            ),
        ]
        self.message = "\n".join(line for line in message_lines if line)
        super().__init__(self.message)

    def _prettify_extra_queries(self, extra_queries):
        line_template = textwrap.dedent(
            """
            {i}. {repeated}
            {stack}
            SQL: {sql}""".lstrip(
                "\n"
            )
        )

        repetitions_dict = defaultdict(list)
        for i, q in enumerate(extra_queries):
            repetitions_dict[q["raw_sql"]].append((i, q))

        deduped_extra_queries_dict = {}
        for repetitions in repetitions_dict.values():
            (i, q) = repetitions[0]
            deduped_extra_queries_dict[i] = q

        line_list = []
        for i, q in deduped_extra_queries_dict.items():
            count = len(repetitions_dict[q["raw_sql"]])
            line = line_template.format(
                i=i,
                repeated=(f"(repeats {count}x, stack only for 1st)" if count > 1 else ""),
                stack=(
                    "\n".join(f"[{s.filename}:{s.lineno}]" for s in q["stack"])
                    if q["stack"]
                    else "(stack unavailable)"
                ),
                sql=q["sql"],
            )
            line_list.append(line)

        pretty = "\n".join(line_list)
        pretty = "\n\t".join(s for s in pretty.split("\n") if s)  # strip empty lines, then ident
        return "\n\t" + pretty


class max_query_count(ContextDecorator):  # noqa
    def __init__(
        self,
        *,
        max_queries=0,
        only_count_select=False,
        affected_cls_or_fn_name=None,
        view_cls_name=None,
        request=None,
        only_log=False,
    ):
        self._max_query_count = max_queries
        self._only_count_select = only_count_select
        self._query_count_before = 0
        self._query_count_after = 0
        self._affected_cls_or_fn_name = affected_cls_or_fn_name
        self._view_cls_name = view_cls_name
        self._request = request
        self._only_log = only_log

    @classmethod
    def for_serializer(cls, serializer_instance, **kwargs):
        auto_kwargs = dict(
            affected_cls_or_fn_name=serializer_instance.__class__.__name__,
            view_cls_name=(
                serializer_instance.context.get("view").__class__.__name__
                if serializer_instance.context.get("view")
                else None
            ),
            request=serializer_instance.context.get("request"),
        )
        auto_kwargs.update(kwargs)
        return cls(**auto_kwargs)

    def get_queries(self):
        if settings.DEBUG:
            # ignore EXPLAIN queries generated by debug tools like django-silk
            queries = [
                q
                for q in self._native_query_capture.captured_queries
                if not q["sql"].strip().startswith("EXPLAIN")
            ]
        else:
            queries = self._native_query_capture.captured_queries
        if self._only_count_select:
            queries = [q for q in queries if q["sql"].strip().startswith("SELECT")]
        return queries

    def __enter__(self):
        self._exit_stack = ExitStack().__enter__()
        self._native_query_capture = native_query_capture()
        self._exit_stack.enter_context(self._native_query_capture)

        self._query_count_before = len(self.get_queries())
        return self

    def __exit__(self, *args):
        self._exit_stack.close()

        queries = self.get_queries()
        self._query_count_after = len(queries)
        actual_query_count = self._query_count_after - self._query_count_before
        extra_queries = queries[-actual_query_count:]

        self._forbid_unoptimized_query(actual_query_count, self._max_query_count, extra_queries)

    def _forbid_unoptimized_query(self, actual_query_count, max_queries, extra_queries):
        if self._has_unoptimized_query(actual_query_count, max_queries):
            exc = QueryCountExceededException(
                expected_query_count=self._max_query_count,
                actual_query_count=actual_query_count,
                extra_queries=extra_queries,
                affected_cls_or_fn_name=self._affected_cls_or_fn_name,
                request=self._request,
                view_cls_name=self._view_cls_name,
            )
            if not self._should_raise_exception():
                logger.warning(exc.message)
            else:
                raise exc

    # These 2 methods where extracted into methods to make them easier to mock in the tests
    def _has_unoptimized_query(self, actual_query_count, max_queries):
        return actual_query_count > max_queries

    def _should_raise_exception(self):
        return settings.DEBUG and not self._only_log
