import keyword
import os
import sys
import token
import tokenize
import warnings
from io import StringIO
from typing import TYPE_CHECKING, Any, TypeAlias

import pygments
from pygments.style import Style
from pygments.token import Token, _TokenType
from functools import cache

from typing import TypedDict

if TYPE_CHECKING:
    # importing this drags in the whole `pygments.formatters` package, which
    # is only needed once something is actually formatted -- see
    # Theme._get_formatter below.
    from pygments.formatters.terminal256 import Terminal256Formatter


TokenStream: TypeAlias = list[tuple[_TokenType, str]]


__all__ = ["Parser", "Theme"]


# Which ``pygments/styles/*.py`` module (and class in it) defines each of the
# builtin pygments styles the themes IPython ships use as a `Theme.base`.
#
# This is only a shortcut, never the source of truth: any name missing from
# here, and any entry that no longer resolves, falls back to pygments' own
# `get_style_by_name`, plugins and all. See `_pygments_base_styles`.
_BUILTIN_PYGMENTS_STYLES: dict[str, tuple[str, str]] = {
    "default": ("default", "DefaultStyle"),
    "gruvbox-dark": ("gruvbox", "GruvboxDarkStyle"),
    "monokai": ("monokai", "MonokaiStyle"),
    "pastie": ("pastie", "PastieStyle"),
}


def _exec_pygments_style_module(module: str, class_name: str) -> Any | None:
    """Read one style's ``styles`` mapping out of ``pygments/styles/<module>.py``.

    Returns None if pygments is not laid out as expected, leaving it to the
    caller to fall back to `pygments.styles.get_style_by_name`.

    The module is executed in isolation and deliberately *not* registered in
    `sys.modules`: registering a submodule of a package that is not itself
    imported breaks later ``import pygments.styles.<module>`` statements, and
    keeping `pygments.styles` unimported is the entire point. Style modules
    are pure data, so executing one twice is harmless, and nothing but the
    ``styles`` dict escapes this function.
    """
    from importlib.machinery import PathFinder
    from importlib.util import module_from_spec

    package = PathFinder.find_spec("pygments.styles", list(pygments.__path__))
    if package is None or package.submodule_search_locations is None:
        return None
    spec = PathFinder.find_spec(
        f"pygments.styles.{module}", list(package.submodule_search_locations)
    )
    if spec is None or spec.loader is None:
        return None
    style_module = module_from_spec(spec)
    try:
        spec.loader.exec_module(style_module)
        return getattr(style_module, class_name).styles
    except Exception:
        return None


def _pygments_base_styles(name: str) -> Any:
    """Return the token -> style-string mapping of a pygments style, by name.

    Equivalent to ``pygments.styles.get_style_by_name(name).styles``, which is
    all IPython ever wants from a base style, but able to answer for the
    handful of builtin styles IPython's own themes are based on without
    importing `pygments.styles`.

    Importing that package -- which importing any of its submodules does too --
    runs `pygments.plugin`, and with it `importlib.metadata` and `email`:
    roughly 10ms whose only purpose is to make third party *style plugins*
    findable by name. No theme IPython ships needs that, and this is on the
    startup path.
    """
    target = _BUILTIN_PYGMENTS_STYLES.get(name)
    if target is not None:
        styles = _exec_pygments_style_module(*target)
        if styles is not None:
            return styles

    from pygments.styles import get_style_by_name

    return get_style_by_name(name).styles


class Symbols(TypedDict):
    top_line: str
    arrow_body: str
    arrow_head: str


_default_symbols: Symbols = Symbols(
    top_line="-",
    arrow_body="-",
    arrow_head=">",
)


class Theme:
    name: str
    base: str | None
    extra_style: dict[_TokenType, str]
    symbols: Symbols

    def __init__(
        self,
        name: str,
        base: str | None,
        extra_style: dict[_TokenType, str],
        *,
        symbols: Symbols | None = None,
    ) -> None:
        self.name = name
        self.base = base
        self.extra_style = extra_style
        s: Symbols = symbols if symbols is not None else _default_symbols
        self.symbols = {**_default_symbols, **s}

    @cache
    def as_pygments_style(self) -> type[Style]:
        if self.base is not None:
            base_styles = _pygments_base_styles(self.base)
        else:
            base_styles = {}

        class MyStyle(Style):
            styles = {**base_styles, **self.extra_style}

        return MyStyle

    @cache
    def _get_formatter(self) -> "Terminal256Formatter":
        from pygments.formatters.terminal256 import Terminal256Formatter

        return Terminal256Formatter(style=self.as_pygments_style())

    def format(self, stream: TokenStream) -> str:
        return pygments.format(stream, self._get_formatter())

    def make_arrow(self, width: int) -> str:
        """generate the leading arrow in front of traceback or debugger"""
        if width >= 2:
            return (
                self.symbols["arrow_body"] * (width - 2)
                + self.symbols["arrow_head"]
                + " "
            )
        elif width == 1:
            return self.symbols["arrow_head"]
        return ""


generate_tokens = tokenize.generate_tokens


#############################################################################
### Python Source Parser (does Highlighting)
#############################################################################

_KEYWORD = token.NT_OFFSET + 1
_TEXT = token.NT_OFFSET + 2

# ****************************************************************************

_pygment_token_mapping: dict[int, _TokenType] = {
    token.NUMBER: Token.Literal.Number,
    token.OP: Token.Operator,
    token.STRING: Token.Literal.String,
    token.COMMENT: Token.Comment,
    token.NAME: Token.Name,
    token.ERRORTOKEN: Token.Error,
    _KEYWORD: Token.Keyword,
    _TEXT: Token.Text,
}

# technically BW is not nocolor, we should have a no-style, style
nocolors_theme = Theme("nocolor", None, {})


linux_theme = Theme(
    "linux",
    "monokai",
    {
        Token.Header: "ansibrightred",
        Token.LinenoEm: "ansibrightgreen",
        Token.Lineno: "ansigreen",
        Token.ValEm: "ansibrightblue",
        Token.VName: "ansicyan",
        Token.Caret: "",
        Token.Filename: "ansibrightgreen",
        Token.ExcName: "ansibrightred",
        Token.Topline: "ansibrightred",
        Token.FilenameEm: "ansigreen",
        Token.Normal: "",
        Token.NormalEm: "ansibrightcyan",
        Token.Line: "ansiyellow",
        Token.TB.Name: "ansimagenta",
        Token.TB.NameEm: "ansibrightmagenta",
        Token.Breakpoint: "",
        Token.Breakpoint.Enabled: "ansibrightred",
        Token.Breakpoint.Disabled: "ansired",
        Token.Prompt: "ansibrightgreen",
        Token.PromptNum: "ansigreen bold",
        Token.OutPrompt: "ansibrightred",
        Token.OutPromptNum: "ansired bold",
        Token.TbHighlight: "ansiblack bg:ansiyellow",
    },
)

neutral_pygments_equiv = {
    Token.Header: "ansired",
    Token.LinenoEm: "ansigreen",
    Token.Lineno: "ansibrightgreen",
    Token.ValEm: "ansiblue",
    Token.VName: "ansicyan",
    Token.Caret: "",
    Token.Filename: "ansibrightgreen",
    Token.FilenameEm: "ansigreen",
    Token.ExcName: "ansired",
    Token.Topline: "ansired",
    Token.Normal: "",
    Token.NormalEm: "ansicyan",
    Token.Line: "ansired",
    Token.TB.Name: "ansibrightmagenta",
    Token.TB.NameEm: "ansimagenta",
    Token.Breakpoint: "",
    Token.Breakpoint.Enabled: "ansibrightred",
    Token.Breakpoint.Disabled: "ansired",
    ## specific override of pygments defaults for visibility
    Token.Number: "ansigreen",
    Token.Operator: "noinherit",
    Token.String: "ansiyellow",
    Token.Name.Function: "ansiblue",
    Token.Name.Class: "bold ansiblue",
    Token.Name.Namespace: "bold ansiblue",
    Token.Name.Variable.Magic: "ansiblue",
    Token.Prompt: "ansigreen",
    Token.OutPrompt: "ansired",
    Token.TbHighlight: "ansiblack bg:ansiyellow",
}


neutral_pygments_nt = {
    **neutral_pygments_equiv,
    Token.PromptNum: "ansigreen bold",
    Token.OutPromptNum: "ansired bold",
}
neutral_pygments_posix = {
    **neutral_pygments_equiv,
    Token.PromptNum: "ansibrightgreen bold",
    Token.OutPromptNum: "ansibrightred bold",
}


neutral_nt = Theme("neutral:nt", "default", neutral_pygments_nt)
neutral_posix = Theme("neutral:posix", "default", neutral_pygments_posix)


# Hack: the 'neutral' colours are not very visible on a dark background on
# Windows. Since Windows command prompts have a dark background by default, and
# relatively few users are likely to alter that, we will use the 'Linux' colours,
# designed for a dark background, as the default on Windows. Changing it here
# avoids affecting the prompt colours rendered by prompt_toolkit, where the
# neutral defaults do work OK.
if os.name == "nt":
    neutral_theme = neutral_nt
else:
    neutral_theme = neutral_posix


lightbg_theme = Theme(
    "lightbg",
    "pastie",
    {
        Token.Header: "ansired",
        Token.LinenoEm: "ansigreen",
        Token.Lineno: "ansibrightgreen",
        Token.ValEm: "ansiblue",
        Token.VName: "ansicyan",
        Token.Caret: "",
        Token.Filename: "ansigreen",
        Token.FilenameEm: "ansibrightgreen",
        Token.ExcName: "ansired",
        Token.Topline: "ansired",
        Token.Normal: "",
        Token.NormalEm: "ansicyan",
        Token.Line: "ansired",
        Token.TB.Name: "ansibrightmagenta",
        Token.TB.NameEm: "ansimagenta",
        Token.Breakpoint: "",
        Token.Breakpoint.Enabled: "ansibrightred",
        Token.Breakpoint.Disabled: "ansired",
        Token.Prompt: "ansibrightblue",
        Token.PromptNum: "ansiblue bold",
        Token.OutPrompt: "ansibrightred",
        Token.OutPromptNum: "ansired bold",
        Token.TbHighlight: "ansired bg:ansiyellow",
    },
)

PRIDE_RED = "#E40303"
PRIDE_ORANGE = "#FF8C00"
PRIDE_YELLOW = "#FFED00"
PRIDE_GREEN = "#008026"
PRIDE_INDIGO = "#004CFF"
PRIDE_VIOLET = "#732982"
pride_theme = Theme(
    "pride",
    "pastie",
    {
        Token.Header: PRIDE_INDIGO,
        Token.LinenoEm: f"{PRIDE_GREEN} italic",
        Token.Lineno: f"{PRIDE_GREEN} bold",
        Token.ValEm: f"{PRIDE_INDIGO} italic",
        Token.VName: "ansicyan",
        Token.Caret: "",
        Token.Filename: f"{PRIDE_YELLOW}",
        Token.FilenameEm: f"bg:{PRIDE_VIOLET}",
        Token.ExcName: f"{PRIDE_ORANGE}",
        Token.Topline: f"{PRIDE_RED}",
        Token.Normal: "",
        Token.NormalEm: "bold",
        Token.Line: "ansired",
        Token.TB.Name: "ansibrightmagenta",
        Token.TB.NameEm: "ansimagenta",
        Token.Breakpoint: "",
        Token.Breakpoint.Enabled: "ansibrightred",
        Token.Breakpoint.Disabled: "ansired",
        Token.Prompt: "ansibrightblue",
        Token.Prompt.Continuation.L1: f"ansiwhite bg:{PRIDE_RED}",
        Token.Prompt.Continuation.L2: f"ansiwhite bg:{PRIDE_ORANGE}",
        Token.Prompt.Continuation.L3: f"ansiblack bg:{PRIDE_YELLOW}",
        Token.Prompt.Continuation.L4: f"ansiwhite bg:{PRIDE_GREEN}",
        Token.Prompt.Continuation.L5: f"ansiwhite bg:{PRIDE_INDIGO}",
        Token.Prompt.Continuation.L6: f"ansiwhite bg:{PRIDE_VIOLET}",
        Token.PromptNum: "ansiblue bold",
        Token.OutPrompt: "ansibrightred",
        Token.OutPromptNum: "ansired bold",
        Token.TbHighlight: f"bg:{PRIDE_YELLOW}",
    },
    symbols={"arrow_body": "\u2500", "arrow_head": "\u25b6", "top_line": "\u2500"},
)


C1 = "#D52D00"
C2 = "#EF7627"
C3 = "#FF9A56"
White = "#FFFFFF"
C5 = "#D162A4"
C6 = "#B55690"
C7 = "#A30262"

pl = {
    # Token.Whitespace: "#bbbbbb",
    Token.Comment: "#888888",
    Token.String: C5,
    Token.String.Escape: C1,
    Token.Keyword: f"italic {C2}",
    Token.Name.Class: C2,
    Token.Name.Exception: C1,
    Token.Name.Builtin: C3,
    Token.Name.Variable: C6,
    Token.Name.Constant: C7,
    Token.Name.Decorator: C2,
    Token.Number: C7,
    Token.Generic.Deleted: f"bg:{C1} #000000",
    Token.Generic.Emph: "italic",
    Token.Generic.Strong: "bold",
    Token.Generic.EmphStrong: "bold italic",
}

pridel_theme = Theme(
    "pride:l",
    None,
    {
        Token.Header: C3,
        Token.LinenoEm: C3,
        Token.Lineno: C2,
        Token.ValEm: C2,
        Token.VName: C2,
        Token.Caret: "",
        Token.Filename: C2,
        Token.FilenameEm: C3,
        Token.ExcName: C1,
        Token.Topline: C1,
        Token.Normal: "",
        Token.NormalEm: "bold",
        Token.Line: C2,
        Token.TB.Name: C6,
        Token.TB.NameEm: C7,
        Token.Breakpoint: "",
        Token.Breakpoint.Enabled: C1,
        Token.Breakpoint.Disabled: C7,
        Token.Prompt: C1,
        Token.PromptNum: C2,
        Token.Prompt.Continuation: C7,
        Token.Prompt.Continuation.L1: C2,
        Token.Prompt.Continuation.L2: C3,
        Token.Prompt.Continuation.L3: White,
        Token.Prompt.Continuation.L4: C5,
        Token.Prompt.Continuation.L5: C6,
        Token.Prompt.Continuation.L6: C7,
        Token.OutPrompt: C6,
        Token.OutPromptNum: C5,
        **pl,
    },
    symbols={"arrow_body": "\u2500", "arrow_head": "\u25b6", "top_line": "\u2500"},
)

GRUVBOX_VAL_EM = "#D79921"
GRUVBOX_V_NAME = "#83A598"
GRUVBOX_FILENAME = "#FBF1C7"
GRUVBOX_EXCEPTION_NAME = "#FB4934"
GRUVBOX_TOPLINE = "#CC241D"
GRUVBOX_BREAKPOINT_ENABLED = "#FB4934"
GRUVBOX_BREAKPOINT_DISABLED = "#CC241D"
GRUVBOX_PROMPT = "#689D6A"
GRUVBOX_PROMPT_NUM = "#8EC07C"
GRUVBOX_OUT_PROMPT = "#B16286"
GRUVBOX_OUT_PROMPT_NUM = "#D3869B"
gruvbox_dark_theme = Theme(
    "gruvbox-dark",
    "gruvbox-dark",
    {
        Token.Lineno: GRUVBOX_PROMPT_NUM,
        Token.LinenoEm: f"{GRUVBOX_PROMPT_NUM} bold",
        Token.ValEm: f"{GRUVBOX_VAL_EM} bold",
        Token.VName: GRUVBOX_V_NAME,
        Token.Caret: "",
        Token.Filename: GRUVBOX_FILENAME,
        Token.FilenameEm: f"{GRUVBOX_FILENAME} bold",
        Token.ExcName: f"{GRUVBOX_EXCEPTION_NAME} bold",
        Token.Topline: GRUVBOX_TOPLINE,
        Token.Breakpoint.Enabled: GRUVBOX_BREAKPOINT_ENABLED,
        Token.Breakpoint.Disabled: GRUVBOX_BREAKPOINT_DISABLED,
        Token.Prompt: GRUVBOX_PROMPT,
        Token.PromptNum: f"{GRUVBOX_PROMPT_NUM} bold",
        Token.OutPrompt: GRUVBOX_OUT_PROMPT,
        Token.OutPromptNum: f"{GRUVBOX_OUT_PROMPT_NUM} bold",
        Token.TbHighlight: f"bg:{GRUVBOX_TOPLINE}",
    },
    symbols={"arrow_body": "\u2500", "arrow_head": "\u25b6", "top_line": "\u2500"},
)

theme_table: dict[str, Theme] = {
    "nocolor": nocolors_theme,
    "linux": linux_theme,
    "neutral": neutral_theme,
    "neutral:nt": neutral_nt,
    "neutral:posix": neutral_posix,
    "lightbg": lightbg_theme,
    "pride": pride_theme,
    "pride:l": pridel_theme,
    "gruvbox-dark": gruvbox_dark_theme,
}


class Parser:
    """Format colored Python source."""

    _theme_name: str
    out: Any
    pos: int
    lines: list[int]
    raw: str

    def __init__(self, out: Any = sys.stdout, *, theme_name: str | None = None) -> None:
        """Create a parser with a specified color table and output channel.

        Call format() to process code.
        """

        assert theme_name is not None

        self.out = out
        self.pos = 0
        self.lines = []
        self.raw = ""
        if theme_name is not None:
            if theme_name in ["Linux", "LightBG", "Neutral", "NoColor"]:
                warnings.warn(
                    f"Theme names and color schemes are lowercase in IPython 9.0 use {theme_name.lower()} instead",
                    DeprecationWarning,
                    stacklevel=2,
                )
                theme_name = theme_name.lower()
        if not theme_name:
            self.theme_name = "nocolor"
        else:
            self.theme_name = theme_name

    @property
    def theme_name(self) -> str:
        return self._theme_name

    @theme_name.setter
    def theme_name(self, value: str) -> None:
        assert value == value.lower()
        self._theme_name = value

    @property
    def style(self) -> str:
        # `style` was renamed `theme_name`; this always raises to catch
        # leftover callers of the old name. Body kept only to satisfy the -> str
        # return type.
        assert False
        return self._theme_name  # type: ignore[unreachable]

    @style.setter
    def style(self, val: str) -> None:
        assert False
        assert val == val.lower()  # type: ignore[unreachable]
        self._theme_name = val

    def format(self, raw: str, out: Any = None) -> str | None:
        return self.format2(raw, out)[0]

    def format2(self, raw: str, out: Any = None) -> tuple[str | None, bool]:
        """Parse and send the colored source.

        If out is not specified, the defaults (given to constructor) are used.

        out should be a file-type object. Optionally, out can be given as the
        string 'str' and the parser will automatically return the output in a
        string."""

        string_output = 0
        if out == "str" or self.out == "str" or isinstance(self.out, StringIO):
            # XXX - I don't really like this state handling logic, but at this
            # point I don't want to make major changes, so adding the
            # isinstance() check is the simplest I can do to ensure correct
            # behavior.
            out_old = self.out
            self.out = StringIO()
            string_output = 1
        elif out is not None:
            self.out = out
        else:
            raise ValueError(
                '`out` or `self.out` should be file-like or the value `"str"`'
            )

        # Fast return of the unmodified input for nocolor scheme
        # TODO:
        if self.theme_name == "nocolor":
            error = False
            self.out.write(raw)
            if string_output:
                return raw, error
            return None, error

        # local shorthands

        # Remove trailing whitespace and normalize tabs
        self.raw = raw.expandtabs().rstrip()

        # store line offsets in self.lines
        self.lines = [0, 0]
        pos = 0
        raw_find = self.raw.find
        lines_append = self.lines.append
        while True:
            pos = raw_find("\n", pos) + 1
            if not pos:
                break
            lines_append(pos)
        lines_append(len(self.raw))

        # parse the source and write it
        self.pos = 0
        text = StringIO(self.raw)

        error = False
        try:
            for atoken in generate_tokens(text.readline):
                self(*atoken)
        except tokenize.TokenError as ex:
            msg = ex.args[0]
            line = ex.args[1][0]
            self.out.write(
                theme_table[self.theme_name].format(
                    [
                        (Token, "\n\n"),
                        (
                            Token.Error,
                            f"*** ERROR: {msg}{self.raw[self.lines[line] :]}",
                        ),
                        (Token, "\n"),
                    ]
                )
            )
            error = True
        self.out.write(
            theme_table[self.theme_name].format(
                [
                    (Token, "\n"),
                ]
            )
        )

        if string_output:
            output = self.out.getvalue()
            self.out = out_old
            return (output, error)
        return (None, error)

    def _inner_call_(
        self, toktype: int, toktext: str, start_pos: tuple[int, int]
    ) -> str:
        """like call but write to a temporary buffer"""
        srow, scol = start_pos

        # calculate new positions
        oldpos = self.pos
        newpos = self.lines[srow] + scol
        self.pos = newpos + len(toktext)

        # send the original whitespace, if needed
        if newpos > oldpos:
            acc = self.raw[oldpos:newpos]
        else:
            acc = ""

        # skip indenting tokens
        if toktype in [token.INDENT, token.DEDENT]:
            self.pos = newpos
            return acc

        # map token type to a color group
        if token.LPAR <= toktype <= token.OP:
            toktype = token.OP
        elif toktype == token.NAME and keyword.iskeyword(toktext):
            toktype = _KEYWORD
        pyg_tok_type = _pygment_token_mapping.get(toktype, Token.Text)

        # send text, pygments should take care of splitting on newline and resending
        # the correct self.colors after the new line, which is necessary for pagers
        acc += theme_table[self.theme_name].format([(pyg_tok_type, toktext)])
        return acc

    def __call__(
        self,
        toktype: int,
        toktext: str,
        start_pos: tuple[int, int],
        end_pos: tuple[int, int],
        line: str,
    ) -> None:
        """Token handler, with syntax highlighting."""
        self.out.write(self._inner_call_(toktype, toktext, start_pos))
