from __future__ import annotations

import dataclasses
import logging
from collections.abc import Collection, Mapping, Sequence
from contextvars import ContextVar, Token
from copy import copy
from dataclasses import dataclass, field
from types import TracebackType
from typing import ClassVar, Optional, Self

from .exception import DiagnosticException, T
from .models import (
    DiagnosticError,
    Loc,
    LocList,
    LocPattern,
    LocReplacement,
    PlainLocPattern,
    PlainLocWithReplacement,
    loc_from_plain,
    loc_from_plain_pattern,
    loc_from_plain_replacement,
)


@dataclass
class _DiagnosticContext:
    """Holds current context of diagnostic collector."""

    self: DiagnosticContext | None = None
    """DiagnosticContext that created this context."""

    prefix: LocList = field(default_factory=LocList)
    """Holds prefix for the errors within this context."""

    suffix: LocList = field(default_factory=LocList)
    """Holds suffix for the errors within this context."""

    strip_prefixes: list[LocPattern] = field(default_factory=list)
    """List of prefixes to be stripped from the errors in this context."""

    strip_suffixes: list[LocPattern] = field(default_factory=list)
    """List of suffixes to be stripped from the errors in this context."""

    mapping: dict[LocPattern, LocReplacement] = field(default_factory=dict)
    """Mapping to translate specific locations to other locations."""

    parent: _DiagnosticContext | None = None
    """Parent context of this context, if any. It is used to combine prefixes and suffixes to provide full location."""

    def apply(self, loc: LocList) -> LocList:
        """
        Apply prefix and suffix to the location.
        """
        out: LocList = LocList(loc)

        for prefix in self.strip_prefixes:
            if prefix.match(out[0 : len(prefix)]) is not False:
                out = LocList(out[len(prefix) :])
                break

        for suffix in self.strip_suffixes:
            if suffix.match(out[-len(suffix) :]) is not False:
                out = LocList(out[: -len(suffix)])
                break

        out = LocList([*self.prefix, *out, *self.suffix])

        for pattern, replacement in self.mapping.items():
            matches = pattern.match(out)
            if matches is not False:
                out = replacement.apply(matches)
                break

        if self.parent:
            out = self.parent.apply(out)

        return out

    def __repr__(self) -> str:
        items: list[str] = []

        if self.prefix:
            items.append(f"prefix={self.prefix}")

        if self.suffix:
            items.append(f"suffix={self.suffix}")

        if self.strip_prefixes:
            items.append(f"strip_prefixes={self.strip_prefixes}")

        if self.strip_suffixes:
            items.append(f"strip_suffixes={self.strip_suffixes}")

        if self.parent:
            items.append(f"parent={self.parent!r}")

        if items:
            return f"<DiagnosticContext {', '.join(items)}>"

        return "<DiagnosticContext>"


class DiagnosticContext:
    """
    Serves as unified location context for subsequent diagnostics. Can be used as context manager or manually, but must be
    always used only within DiagnosticCollector context.
    """

    _collector_context: ClassVar[ContextVar[DiagnosticCollector | None]] = ContextVar[Optional["DiagnosticCollector"]](
        "DiagnosticContext._collector_context", default=None
    )
    _diagnostic_context: ClassVar[ContextVar[_DiagnosticContext]] = ContextVar[_DiagnosticContext](
        "DiagnosticContext._diagnostic_context",
        default=_DiagnosticContext(),  # noqa: B039
    )

    def __init__(
        self,
        *,
        prefix: LocList | Loc | None = None,
        strip_prefix: LocList | Loc | LocPattern | None = None,
        strip_prefixes: Collection[LocList | Loc | LocPattern] | None = None,
        suffix: LocList | Loc | None = None,
        strip_suffix: LocList | Loc | LocPattern | None = None,
        strip_suffixes: Collection[LocList | Loc | LocPattern] | None = None,
        mapping: Mapping[LocList | Loc | PlainLocPattern | LocPattern, LocList | LocReplacement | PlainLocWithReplacement | Loc]
        | None = None,
    ) -> None:
        """
        Args:
            prefix: Add specified prefix to all errors in this collector.
            strip_prefix: Strip specified prefix from all errors in this collector, including errors from inner collectors.
            strip_prefixes: Strip specified prefixes from all errors in this collector, including errors from inner collectors
                (using this argument, you can specify multiple different prefixes to be stripped).
            suffix: Add specified suffix to all errors in this collector.
            strip_suffix: Strip specified suffix from all errors in this collector, including errors from inner collectors.
            strip_suffixes: Strip specified suffixes from all errors in this collector, including errors from inner collectors
                (using this argument, you can specify multiple different suffixes to be stripped).
            mapping: Mapping to translate specific locations to other locations. This is applied after stripping prefixes and
                suffixes and after applying this context's own prefix and suffix, to allow completely rewriting certain locations
                to anything arbitrary without being restrained by the prefix and suffix of this context. The mapping is searched
                in the order it has been defined. After successful match, the rest of mapping is ignored for that particular
                error.

        Raises:
            TypeError: When prefix or suffix is not a list of locations.
        """
        if not isinstance(prefix, list) and prefix is not None:
            raise TypeError("Prefix must be a list of locations.")

        if not isinstance(suffix, list) and suffix is not None:
            raise TypeError("Suffix must be a list of locations.")

        self._init_context = _DiagnosticContext(
            prefix=loc_from_plain(prefix or []),
            suffix=loc_from_plain(suffix or []),
            strip_prefixes=([loc_from_plain_pattern(strip_prefix)] if strip_prefix else [])
            + list(map(loc_from_plain_pattern, strip_prefixes or [])),
            strip_suffixes=([loc_from_plain_pattern(strip_suffix)] if strip_suffix else [])
            + list(map(loc_from_plain_pattern, strip_suffixes or [])),
            mapping={loc_from_plain_pattern(k): loc_from_plain_replacement(v) for k, v in (mapping or {}).items()},
        )

        self._diagnostic_context_token: Token[_DiagnosticContext] | None = None

    @property
    def collector(self) -> DiagnosticCollector:
        """
        Access the DiagnosticCollector this DiagnosticContext appends errors to. It is the first collector in the hierarchy
        of contexts.

        Exception is raised when used outside DiagnosticCollector context.

        Raises:
            AttributeError: When used outside DiagnosticCollector context.
        """
        collector = DiagnosticContext._collector_context.get()

        # Test explicitely for None, as DiagnosticCollector contains __bool__ method which checks for presence of errors.
        if collector is None:
            raise AttributeError("DiagnosticContext must be used within DiagnosticCollector context.")

        return collector

    @property
    def context(self) -> _DiagnosticContext:
        """
        Access the current context of the diagnostic collector.
        """
        return DiagnosticContext._diagnostic_context.get()

    def _append(self, error: DiagnosticError, prefix: LocList | Loc | None = None, suffix: LocList | Loc | None = None) -> None:
        """
        Append error to diagnostics.

        Args:
            error: Error to be appended.
            prefix: Prefix to prepend to error location.
            suffix: Suffix to append to error location.

        Raises:
            TypeError: When error is not DiagnosticError.
        """
        if not isinstance(error, DiagnosticError):
            raise TypeError("DiagnosticCollector can accept only DiagnosticErrors.")

        ctx = self.context
        ctx = dataclasses.replace(
            ctx,
            prefix=loc_from_plain(prefix) if prefix is not None else ctx.prefix,
            suffix=loc_from_plain(suffix) if suffix is not None else ctx.suffix,
        )

        logging.debug("Error context: %r", ctx)

        error = error.model_copy(update={"loc": ctx.apply(loc_from_plain(error.loc))})
        self.collector.errors.append(error)

    def append(
        self,
        error: DiagnosticError,
        prefix: LocList | Loc | None = None,
        suffix: LocList | Loc | None = None,
        stacklevel: int = 0,
    ) -> Self:
        """
        Append new error to the collector, optionally prefixing it with specified prefix or suffixing it with specified suffix.

        Args:
            error: Error to be added.
            prefix: Prefix to be prepended before error's loc. If specified, overwrites class-level prefix.
            suffix: Suffix to be appended to error's loc. If specified, overwrites class-level suffix.
            stacklevel: Optional stacklevel for logging (default 0).

        Returns:
            Self for chaining
        """
        # Skip logging + self.append() + parent stacklevel => log first stack frame outside diagnostic library,
        # to provide correct context for the user.
        stacklevel += 2

        if prefix and suffix:
            self.collector.logger.debug(
                "Appending error with prefix %s and suffix %s: %s", prefix, suffix, error, stacklevel=stacklevel
            )
        elif prefix:
            self.collector.logger.debug("Appending error with prefix %s: %s", prefix, error, stacklevel=stacklevel)
        elif suffix:
            self.collector.logger.debug("Appending error with suffix %s: %s", suffix, error, stacklevel=stacklevel)
        else:
            self.collector.logger.debug("Appending error: %s", error, stacklevel=stacklevel)

        self._append(error, prefix, suffix)

        return self

    def add(
        self,
        error: DiagnosticError,
        prefix: LocList | Loc | None = None,
        suffix: LocList | Loc | None = None,
        stacklevel: int = 0,
    ) -> Self:
        """Alias for append()"""
        return self.append(error, prefix, suffix, stacklevel + 1)

    def include(
        self,
        other: DiagnosticCollector | DiagnosticException[T] | Sequence[DiagnosticError],
        prefix: LocList | Loc | None = None,
        suffix: LocList | Loc | None = None,
        stacklevel: int = 0,
    ) -> Self:
        """
        Include errors from other diagnostic response, optionally prefixing errors with location.

        Args:
            other: Other diagnostic to include.
            prefix: Error prefix. Overrides class-level prefix.
            suffix: Error suffix. Overrides class-level suffix.
            stacklevel: Optional stacklevel for logging (default 0).

        Returns:
            Self for chaining

        Raises:
            TypeError: When other is not DiagnosticCollector, DiagnosticException or sequence of DiagnosticErrors.
        """
        errors: list[DiagnosticError] = []

        if isinstance(other, DiagnosticCollector):
            if other == self:
                return self

            errors.extend(other.errors)

            # As errors are processed by including in this collector, clear errors from other, to avoid raising exception.
            other.errors = []
        elif isinstance(other, DiagnosticException):
            # As the exception contains finished errors with correct loc, append the errors directly to the collector's errors
            # container without modifying the loc with current context, as it would cause doubling prefixes / suffixes.
            self.collector.errors.extend(other.errors())
            return self
        elif not isinstance(other, Sequence):
            raise TypeError(
                "DiagnosticCollector can accept only other DiagnosticCollectior, DiagnosticException or sequence of "
                "DiagnosticErrors."
            )
        else:
            errors.extend(other)

        stacklevel += 2

        for error in errors:
            if prefix is not None and suffix is not None:
                self.collector.logger.debug(
                    "Including error from nested diagnostics with prefix %s and suffix %s: %s",
                    prefix,
                    suffix,
                    error,
                    stacklevel=stacklevel,
                )
            elif prefix is not None:
                self.collector.logger.debug(
                    "Including error from nested diagnostics with prefix %s: %s", prefix, error, stacklevel=stacklevel
                )
            elif suffix is not None:
                self.collector.logger.debug(
                    "Including error from nested diagnostics with suffix %s: %s", suffix, error, stacklevel=stacklevel
                )
            else:
                self.collector.logger.debug("Including error from nested diagnostics: %s", error, stacklevel=stacklevel)

            self._append(error, prefix, suffix)

        return self

    def __enter__(self) -> Self:
        parent_context = self.context

        # DiagnosticCollector overwrites context without parent, as it catches all errors and does not
        # populate it to the upper context. It acts as a root context that raises the exception on exit,
        # which in turn can be catched by higher order collector to include its errors, which in turn
        # completes the locations with higher order collector's context.
        self._diagnostic_context_token = self.__class__._diagnostic_context.set(  # noqa: SLF001
            _DiagnosticContext(
                self=self,
                prefix=self._init_context.prefix,
                suffix=self._init_context.suffix,
                strip_prefixes=self._init_context.strip_prefixes,
                strip_suffixes=self._init_context.strip_suffixes,
                mapping=self._init_context.mapping,
                parent=parent_context,
            )
        )

        self.collector.logger.debug(
            "Entering diagnostic context %r",
            self.__class__._diagnostic_context.get(),  # noqa: SLF001
            # Skip logging + self.__enter__ => log first stack frame outside diagnostic library, to provide
            stacklevel=2,
        )

        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
        stacklevel: int = 0,
    ) -> None:
        self.collector.logger.debug(
            "Leaving diagnostic context %r",
            self.__class__._diagnostic_context.get(),  # noqa: SLF001
            # Skip logging + self.__exit__ + parent stacklevel => log first stack frame outside diagnostic library, to provide
            # correct context for the user.
            stacklevel=2 + stacklevel,
        )

        if self._diagnostic_context_token:
            self._diagnostic_context.reset(self._diagnostic_context_token)

    async def __aenter__(self) -> Self:
        """
        Async version of context manager. See __enter__().
        """
        return self.__enter__()

    async def __aexit__(
        self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
    ) -> None:
        """
        Async version of context manager. See __exit__().
        """
        self.__exit__(exc_type, exc_val, exc_tb)


class DiagnosticCollector(DiagnosticContext):
    """
    Diagnostic collector, that catches errors. Can be used as context manager, then it raises automatically
    at the __exit__ of the with block. Or can be used manually without context manager, then just call
    self.raise_if_errors().

    Example:
    >>> with DiagnosticCollector() as diag:
    >>>     diag.append(DiagnosticError(loc=["somewhere"], msg="There was an error.", type="error"))
    >>> # Here, DiagnosticException is raised.

    Or:

    >>> diag = DiagnosticCollector()
    >>> diag.append(DiagnosticError(loc=["somewhere"], msg="There was an error.", type="error"))
    >>> diag.raise_if_errors()  # Here, DiagnosticException is raised.

    Diagnostic collectors can be nested, and errors from inner collectors are included in outer collector.
    """

    def __init__(
        self,
        *,
        prefix: LocList | Loc | None = None,
        strip_prefix: LocList | Loc | LocPattern | None = None,
        strip_prefixes: Collection[LocList | Loc | LocPattern] | None = None,
        suffix: LocList | Loc | None = None,
        strip_suffix: LocList | Loc | LocPattern | None = None,
        strip_suffixes: Collection[LocList | Loc | LocPattern] | None = None,
        mapping: Mapping[LocList | Loc | PlainLocPattern | LocPattern, LocList | LocReplacement | PlainLocWithReplacement | Loc]
        | None = None,
        logger: logging.Logger | None = None,
    ) -> None:
        """
        Args:
            prefix: Add specified prefix to all errors in this collector.
            strip_prefix: Strip specified prefix from all errors in this collector, including errors from inner collectors.
            strip_prefixes: Strip specified prefixes from all errors in this collector, including errors from inner collectors
                (using this argument, you can specify multiple different prefixes to be stripped).
            suffix: Add specified suffix to all errors in this collector.
            strip_suffix: Strip specified suffix from all errors in this collector, including errors from inner collectors.
            strip_suffixes: Strip specified suffixes from all errors in this collector, including errors from inner collectors
                (using this argument, you can specify multiple different suffixes to be stripped).
            mapping: Mapping to translate specific locations to other locations. This is applied after stripping prefixes and
                suffixes and after applying this context's own prefix and suffix, to allow completely rewriting certain locations
                to anything arbitrary without being restrained by the prefix and suffix of this context. The mapping is searched
                in the order it has been defined. After successful match, the rest of mapping is ignored for that particular
                error.
            logger: Optional logger to use for logging errors. If not specified, uses "diagnostics" logger.
        """
        super().__init__(
            prefix=prefix,
            strip_prefix=strip_prefix,
            strip_prefixes=strip_prefixes,
            suffix=suffix,
            strip_suffix=strip_suffix,
            strip_suffixes=strip_suffixes,
            mapping=mapping,
        )

        # Set to true in raise_if_errors, to prevent doubling of errors when context manager catches own exception.
        self._raised_from_self = False

        self.logger = logger or logging.getLogger("diagnostics")
        self.errors: list[DiagnosticError] = []

        # Set to true in raise_if_errors, to prevent doubling of errors when context manager catches own exception.
        self._raised_from_self = False

        self._collector_context_token: Token[DiagnosticCollector | None] | None = None

    @property
    def collector(self) -> DiagnosticCollector:
        """
        Diagnostic collector collects it's own errors, does not propagate it to the context, which can be any other collector.
        """
        return self

    def raise_if_errors(self, stacklevel: int = 0) -> None:
        """
        Raises:
             DiagnosticException: if there are any collected errors. Otherwise, does nothing.

        Args:
            stacklevel: Stack level for the logging purposes.
        """
        if bool(self):
            self._raised_from_self = True
            # Skip logging + self.raise_if_errors => log first stack frame outside diagnostic library, to provide
            # correct context for the user.
            self.logger.debug("Raising exception from %r: %s", self.context, self.errors, stacklevel=stacklevel + 2)
            raise DiagnosticException(detail=copy(self.errors), logger=self.logger, raised_from=self)

    def __bool__(self) -> bool:
        """
        Whether the exception contains any error, therefore should be raised.
        """
        return bool(self.errors)

    def __enter__(self) -> Self:
        """
        Start context manager. At end of context, DiagnosticException is automatically raised when there are any errors.
        """
        self._collector_context_token = self.__class__._collector_context.set(self)  # noqa: SLF001
        return super().__enter__()

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
        stacklevel: int = 0,
    ) -> None:
        """
        Automatically raises DiagnosticException if there is any diagnostic to be presented.

        If __exit__ occured because of another DiagnosticException, append errors from it to this diagnostic
        (with optionally specified prefix) and re-raise new DiagnosticException with all collected errors.
        """
        # Include errors from inner exception from inner collector.
        if isinstance(exc_val, DiagnosticException) and exc_val.raised_from != self:
            # include errors from inner exception, ensuring prefix and suffix are not modified, as it is already an error
            # with correct context.
            # stacklevel to jump out of __exit__
            self.include(exc_val, stacklevel=stacklevel + 1, prefix=[], suffix=[])

        # stacklevel to jump out of __exit__
        super().__exit__(exc_type, exc_val, exc_tb, stacklevel=stacklevel + 1)  # type: ignore[call-arg]

        if self._collector_context_token:
            self.__class__._collector_context.reset(self._collector_context_token)  # noqa: SLF001

        # Do not raise DiagnosticException if other exception (other than DiagnosticException) was raised,
        # as we don't want to shadow internal exceptions from user even if there are some diagnostics.
        if exc_val is None or isinstance(exc_val, DiagnosticException):
            self.raise_if_errors(stacklevel=stacklevel + 1)
