import dis
import functools
import re
import types
from sys import version_info

import linetable

import kajiki
from kajiki import i18n, lnotab
from kajiki.html_utils import HTML_EMPTY_ATTRS
from kajiki.ir import generate_python
from kajiki.util import flattener, literal


class _obj:  # noqa: N801
    def __init__(self, **kw):
        for k, v in kw.items():
            setattr(self, k, v)


class _Template:
    """Base Class for all compiled Kajiki Templates.

    All kajiki templates created from a :class:`kajiki.ir.TemplateNode` will
    be subclasses of this class.

    As the template body code runs inside ``__main__`` method of this
    class, the instance of this class is always available as ``self``
    inside the template code.

    This class also makes available some global object inside the
    template code itself:

        - ``local`` which is the instance of the template
        - ``defined`` which checks if the given variable is defined
          inside the template scope.
        - ``Markup`` which marks the passed object as markup code and
          prevents escaping for its content.
        - ``__kj__`` which is a special object used by generated code
          providing features like keeping track of py:with stack or
          or the gettext function used to translate text.
    """

    __methods__ = ()
    loader = None
    base_globals = None
    filename = None

    def __init__(self, context=None):
        if context is None:
            context = {}
        self._context = context
        base_globals = self.base_globals or {}
        self.__globals__ = {
            "local": self,
            "self": self,
            "defined": lambda x: x in self.__globals__,
            "literal": literal,
            "Markup": literal,
            "gettext": i18n.gettext,
            "__builtins__": __builtins__,
            "__kj__": kajiki,
        }
        self.__globals__.update(base_globals)
        for k, v in self.__methods__:
            v = v.bind_instance(self)  # noqa: PLW2901
            setattr(self, k, v)
            self.__globals__[k] = v
        self.__kj__ = _obj(
            extend=self._extend,
            push_switch=self._push_switch,
            pop_switch=self._pop_switch,
            case=self._case,
            import_=self._import,
            escape=self._escape,
            gettext=self._gettext,
            render_attrs=self._render_attrs,
            push_with=self._push_with,
            pop_with=self._pop_with,
            collect=self._collect,
        )
        self._switch_stack = []
        self._with_stack = []
        self.__globals__.update(context)
        self.__globals__["_"] = self.__globals__["gettext"]
        self.__globals__["value_of"] = self.__globals__.get

    def __iter__(self):
        """We convert the chunk to string because it can be of any type
        -- after all, the template supports expressions such as ${x+y}.
        Here, ``chunk`` can be the computed expression result.
        """
        for chunk in self.__main__():
            yield str(chunk)

    def render(self):
        """Render the template to a string."""
        return "".join(self)

    def _gettext(self, s):
        """Used by the code generated by the template to translate static text"""
        return self.__globals__["gettext"](s)

    def _push_with(self, locals_, vars):  # noqa: A002
        """Enter a ``py:with`` block.

        When a ``py:with`` block is encountered, previous values
        of the variables assigned inside the ``py:with`` statement are
        pushed on top of a stack by :class:`kajiki.ir.WithNode` so that
        when the node is exited the previous values can be recovered.
        """
        self._with_stack.append([locals_.get(k, ()) for k in vars])

    def _pop_with(self):
        """Exists a ``py:with`` block.

        When a ``py:with`` block is exited the values stack is popped
        and the head returned to :class:`kajiki.ir.WithNode` so that
        it can set any previously existing variable to its old value.
        """
        return self._with_stack.pop()

    def _extend(self, parent):
        """
        Called when a child template extends a parent template
        the first thing it does is asking the loader of the
        child template to load the parent template
        """
        if isinstance(parent, str):
            parent = self.loader.import_(parent)
        p_inst = parent(self._context)
        p_globals = p_inst.__globals__
        # Find overrides
        for k, v in self.__globals__.items():
            if k == "__main__":
                continue
            if not isinstance(v, TplFunc):
                continue
            p_globals[k] = v
        # Find inherited funcs
        for k, v in p_inst.__globals__.items():
            if k == "__main__":
                continue
            if not isinstance(v, TplFunc):
                continue
            if k not in self.__globals__:
                self.__globals__[k] = v
            if not hasattr(self, k):

                def _(k=k):
                    """Capture the 'k' variable in a closure"""

                    def trampoline(*a, **kw):
                        global parent  # noqa: PLW0602
                        return getattr(parent, k)(*a, **kw)

                    return trampoline

                setattr(self, k, TplFunc(_()).bind_instance(self))
        p_globals["child"] = self
        p_globals["local"] = p_inst
        p_globals["self"] = self.__globals__["self"]
        self.__globals__["parent"] = p_inst
        self.__globals__["local"] = self
        return p_inst

    def _push_switch(self, expr):
        """Enter a ``py:switch`` block.

        Pushes provided value on the stack used
        to check ``py:case`` statements against.

        Calling :meth:`._pop_switch` will exit the switch block.
        """
        self._switch_stack.append(expr)

    def _pop_switch(self):
        """Exit current ``py:switch`` block.

        Pops current value from the stack used
        to check ``py:case`` statements against.
        """
        self._switch_stack.pop()

    def _case(self, obj):
        """Check against current ``py:switch`` value."""
        return obj == self._switch_stack[-1]

    def _import(self, name, alias, gbls):
        # Load template as a fragment to avoid extra <DOCTYPE> in included output.
        # Due to loader cache, this has the side effect that if the same
        # template is both included and used as a standalone page
        # it might act as a fragment or not depending on the order it was loaded.
        # But usually templates meant for inclusion are not standalone pages.
        # Also there is no way to set a template as a fragment once loaded.
        # So we can only do it through the loader.
        tpl_cls = self.loader.import_(name, is_fragment=True)
        if alias is None:
            alias = self.loader.default_alias_for(name)
        r = gbls[alias] = tpl_cls(gbls)
        return r

    def _escape(self, value):
        """Returns the given HTML with ampersands, carets and quotes encoded."""
        if value is None or isinstance(value, flattener):
            return value
        if hasattr(value, "__html__"):
            return value.__html__()
        uval = str(value)
        if self._re_escape.search(uval):  # Scan the string before working.
            # stdlib escape() is inconsistent between Python 2 and Python 3.
            # In 3, html.escape() translates the single quote to '&#39;'
            # In 2.6 and 2.7, cgi.escape() does not touch the single quote.
            # Preserve our tests and Kajiki behaviour across Python versions:
            return uval.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;").replace('"', "&quot;")
            # .replace("'", '&#39;'))
            # Above we do NOT escape the single quote; we don't need it because
            # all HTML attributes are double-quoted in our output.
        return uval

    _re_escape = re.compile(r'&|<|>|"')

    def _render_attrs(self, attrs, mode):
        """Render tag attributes in key="value" format.

        A :class:`kajiki.ir.AttrsNode` will generate
        code that in fact leads to this function to generate
        the html for tag attributes.
        """
        if hasattr(attrs, "items"):
            attrs = attrs.items()
        if attrs is not None:
            for k, v in sorted(attrs):
                if k in HTML_EMPTY_ATTRS and v in (True, False):
                    v = k if v else None  # noqa: PLW2901
                if v is None:
                    continue
                if mode.startswith("html") and k in HTML_EMPTY_ATTRS:
                    yield " " + k.lower()
                else:
                    yield f' {k}="{self._escape(v)}"'

    def _collect(self, it):
        result = []
        for part in it:
            if part is None:
                continue
            if isinstance(part, flattener):
                result.append(str(part.accumulate_str()))
            else:
                result.append(str(part))
        if result:
            return "".join(result)
        return None

    @classmethod
    def annotate_lnotab(cls, py_to_tpl):
        for _name, meth in cls.__methods__:
            meth.annotate_lnotab(cls.filename, py_to_tpl, dict(py_to_tpl))

    def defined(self, name):
        """Check if a variable was provided to the template or not"""
        return name in self._context


def Template(ns):  # noqa: N802
    """Creates a :class:`._Template` subclass from an entity with ``exposed`` functions.

    Kajiki uses classes as containers of the exposed functions for convenience,
    but any object that can have the functions as attributes works.

    To be a valid template the original entity must provide at least a ``__main__``
    function::

        class Example:
            @kajiki.expose
            def __main__():
                yield "Hi"


        t = kajiki.Template(Example)
        output = t().render()

        print(output)
        "Hi"
    """
    dct = {}
    methods = dct["__methods__"] = []
    for name in dir(ns):
        value = getattr(ns, name)
        if getattr(value, "exposed", False):
            methods.append((name, TplFunc(getattr(value, "__func__", value))))
    return type(ns.__name__, (_Template,), dct)


def from_ir(ir_node, base_globals=None):
    """Creates a template class from Intermediate Representation TemplateNode.

    This actually creates the class defined by the TemplateNode by executing
    its code and returns a subclass of it.
    The returned class is a subclass of :class:`kajiki.template._Template`.

    It is possible to use `base_globals` to set context values
    or replace default ones
    """
    if base_globals is None:
        base_globals = {}
    py_lines = list(generate_python(ir_node))
    py_text = "\n".join(map(str, py_lines))
    py_linenos = []
    last_lineno = 0
    py_lineno = 1
    for line in py_lines:
        lno = max(last_lineno, line._lineno or 0)  # noqa: SLF001
        for _ in range(str(line).count("\n") + 1):
            py_linenos.append((py_lineno, lno))
            py_lineno += 1
        last_lineno = lno
    dct = {"kajiki": kajiki}
    try:
        exec(py_text, dct)  # noqa: S102
    except (SyntaxError, IndentationError) as e:  # pragma no cover
        raise KajikiSyntaxError(e.msg, py_text, e.filename, e.lineno, e.offset) from e
    tpl = dct["template"]
    tpl.base_globals = base_globals.copy()
    tpl.base_globals.update(dct)
    tpl.py_text = py_text
    tpl.filename = ir_node.filename
    tpl.annotate_lnotab(py_linenos)
    return tpl


class TplFunc:
    """A template function attached to a _Template.

    By default template functions (ie: __main__) depends
    on variables like ``self``, ``local`` and so on which
    are provided by :class:`._Template`.

    This is used by :meth:`.Template` to create a new
    ``_Template`` with the attached functions.
    """

    def __init__(self, func, inst=None):
        self._func = func
        self._inst = inst
        self._bound_func = None

    def bind_instance(self, inst):
        return TplFunc(self._func, inst)

    def __repr__(self):  # pragma no cover
        if self._inst:
            return f"<bound tpl_function {self._func.__name__!r} of {self._inst!r}>"
        return f"<unbound tpl_function {self._func.__name__!r}>"

    def __call__(self, *args, **kwargs):
        if self._bound_func is None:
            self._bound_func = self._bind_globals(self._inst.__globals__)
        return self._bound_func(*args, **kwargs)

    def _bind_globals(self, globals):  # noqa: A002
        """Return a function which has the globals dict set to 'globals'
        and which flattens the result of self._func'.
        """
        func = types.FunctionType(
            self._func.__code__,
            globals,
            self._func.__name__,
            self._func.__defaults__,
            self._func.__closure__,
        )
        return functools.update_wrapper(lambda *a, **kw: flattener(func(*a, **kw)), func)

    def annotate_lnotab(self, filename, py_to_tpl, py_to_tpl_dct):
        if not py_to_tpl:
            return
        code = self._func.__code__

        new_firstlineno = py_to_tpl_dct.get(code.co_firstlineno, 0)
        if version_info >= (3, 11):
            ltable = linetable.parse_linetable(code.co_linetable, code.co_firstlineno)
            new_lnotab = linetable.generate_linetable(
                (
                    (
                        length,
                        py_to_tpl_dct[start_line],
                        py_to_tpl_dct[end_line],
                        None,
                        None,
                    )
                    if start_line
                    else (length, None, None, None, None)
                    for length, start_line, end_line, *_ in ltable
                ),
                firstlineno=new_firstlineno,
            )
        else:
            new_lnotab_numbers = []
            for bc_off, py_lno in dis.findlinestarts(code):
                tpl_lno = py_to_tpl_dct[py_lno]
                new_lnotab_numbers.append((bc_off, tpl_lno))
            if not new_lnotab_numbers:
                return
            new_lnotab = lnotab.lnotab_string(new_lnotab_numbers, new_firstlineno)
        new_code = patch_code_file_lines(code, filename, new_firstlineno, new_lnotab)
        self._func.__code__ = new_code
        return


def patch_code_file_lines(code, filename, firstlineno, lnotab):
    code_args = (
        code.co_argcount,
        code.co_posonlyargcount,
        code.co_kwonlyargcount,
        code.co_nlocals,
        code.co_stacksize,
        code.co_flags,
        code.co_code,
        code.co_consts,
        code.co_names,
        code.co_varnames,
        filename,
        code.co_name,
        code.co_qualname if version_info >= (3, 11) else "REMOVE",
        firstlineno,
        lnotab,  # linetable for >=3.11 and lnotab for <3.11
        code.co_exceptiontable if version_info >= (3, 11) else "REMOVE",
        code.co_freevars,
        code.co_cellvars,
    )
    return types.CodeType(*(arg for arg in code_args if arg != "REMOVE"))


class KajikiTemplateError(Exception):
    def __init__(self, msg, source, filename, linen, coln):
        super().__init__(f"[{filename}:{linen}] {msg}\n{self._get_source_snippet(source, linen)}")
        self.filename = filename
        self.linenum = linen
        self.colnum = coln

    def _get_source_snippet(self, source, lineno):
        lines = source.splitlines()

        # Lines are 1 indexed, account for that.
        lineno -= 1

        parts = []
        for i in range(lineno - 2, lineno + 2):
            if 0 <= i < len(lines):
                parts.append(  # noqa: PERF401
                    "\t {arrow} {src}\n".format(arrow="-->" if i == lineno else "   ", src=lines[i])
                )
        return "".join(parts)


class KajikiSyntaxError(KajikiTemplateError):
    pass
