import re
import collections.abc
import inspect
import typing as t
from itertools import count
import pickle

from .module import Module
from .parsing import parse_type_hints


def template(base_imports, project_imports, queries):
    return render_template(
        """
        # Auto-generated by embrace.codegen
        {% for imp, alias in sorted(bi.items()) %}
        import {{ imp }} as {{ alias }}
        {% end %}

        {% if pi %}
        if __import__("typing").TYPE_CHECKING:
            {% for imp, alias in sorted(pi.items()) %}
            import {{ imp }} as {{ alias }}
            {% end %}
        {% end %}

        _m = {{ bi.embrace_module }}.Module()

        {% for q in queries %}class _{{ q.name }}_Query({{ bi.embrace_query }}.Query[{{q.return_annotation}}, {{ q.scalar_return_annotation }}]):
            def resultset(
                {% capture args %}self, conn: {{ bi.typing }}.Optional[{{ bi.embrace_types }}.Connection], {% for n, h in q.args %}{{ n }}: {{h}}, {% end %}*, params: {{ bi.typing }}.Optional[{{bi.collections_abc}}.Mapping]=None, debug: bool=False, _module = _m{% endcapture %}
            ) -> _Q.ResultSet[{{ q.return_annotation }}, {{ q.scalar_return_annotation }}]:
                return super().resultset(
                    {% capture callargs %}conn, {% for n in q.arg_names %}{{ n }}={{ n }}, {% end %} params=params, debug=debug {% endcapture %}
                )
            def one({{ args }}) -> {{ q.return_annotation }}:
                return super().one({{ callargs }})
            def first({{ args }}) ->  _T.Union[{{ q.return_annotation }}, None]:
                return super().first({{ callargs }})
            def many({{ args }}) -> _C.Iterator[{{ q.return_annotation }}]:
                return super().many({{ callargs }})
            def column({{ args }}) -> _C.Iterator[{{ q.scalar_return_annotation }}]:
                return super().column({{ callargs }})
            def scalar({{ args }}, **_kw) -> {{ q.scalar_return_annotation }}:
                return super().scalar({{ callargs }}, **_kw)

        {% end %}

        {% for q in queries %}
        _m.queries["{{q.name}}"] = {{ q.name }} = _{{ q.name }}_Query()
        {{ q.name }}.__dict__ = {{ bi.pickle }}.loads({{ repr(q.pickle) }})
        {% end %}

        def __getattr__(name):
            return getattr(_m, name)
        """,
        {"bi": base_imports, "pi": project_imports, "i": base_imports | project_imports, "queries": queries},
    )


def generate_module_source(module):
    queries = []

    # Imports required by embrace's internals
    base_imports = ImportDict({
        "pickle": "_P",
        "typing": "_T",
        "collections.abc": "_C",
        "embrace.types": "_ET",
        "embrace.module": "_M",
        "embrace.query": "_Q",
    })
    # Imports imposed by the project, these will be wrapped in an
    # ``if TYPE_CHECKING`` statement to remove the possibility of circular
    # imports
    project_imports = ImportDict({})

    symbol_aliases = symbol_aliaser()
    for q in sorted(module, key=lambda q: q.name or ""):
        q.resolve_includes(module)
        md = q.metadata
        placeholder_info = q.get_placeholder_info()
        if placeholder_info:
            arg_names, placeholder_hints = zip(*placeholder_info)
            arg_hints = symbol_aliases(
                base_imports, project_imports, parse_type_hints(md.imports, placeholder_hints)
            )
        else:
            arg_names = arg_hints = []

        if q.row_mapper:
            row_spec = q.row_mapper.row_spec_after_joins()
        else:
            row_spec = []

        if row_spec:
            mapped = symbol_aliases(base_imports, project_imports, (mo.mapped for mo in row_spec))
            if len(mapped) > 1:
                return_annotation = f"""tuple[{', '.join(f'"{m}"' for m in mapped)}]"""
            else:
                return_annotation = f'"{mapped[0]}"'
            scalar_return_annotation = f'"{mapped[0]}"'
        else:
            return_annotation = "tuple"
            scalar_return_annotation = f"{base_imports.typing}.Any"

        queries.append(
            AttrDict(
                q.__dict__
                | {
                    "args": list(zip(arg_names, arg_hints)),
                    "arg_names": arg_names,
                    "return_annotation": return_annotation,
                    "scalar_return_annotation": scalar_return_annotation,
                    "pickle": pickle.dumps(q.__dict__),
                }
            )
        )

    return template(base_imports, project_imports, queries)


def symbol_aliaser() -> (
    t.Callable[
        [
            collections.abc.Mapping[str, t.Optional[str]],
            collections.abc.MutableMapping[str, t.Optional[str]],
            t.Iterator
        ], list[str]
    ]
):
    c = count()

    def symbol_aliases(
        base_aliases: collections.abc.Mapping[str, t.Optional[str]],
        project_aliases: collections.abc.Mapping[str, t.Optional[str]],
        obs: t.Iterator[t.Any],
    ) -> list[str]:

        symbols = []
        for ix, ob in enumerate(obs):
            if ob is None:
                symbols.append(f"{base_aliases['typing']}.Any")
                continue
            try:
                ra = t.get_origin(inspect.signature(ob).return_annotation)
                if ra is not inspect._empty and ra is not None:
                    ob = ra
            except (TypeError, ValueError):
                pass
            try:
                mod = ob.__module__
            except AttributeError:
                symbols.append(f"{base_aliases['typing']}.Any")
            else:
                name = getattr(ob, "__qualname__", None)
                if name is None or "<locals>" in name or mod == "embrace.mapobject":
                    # <locals>: non-exportable symbol,
                    # typically the result of calling mapobject.dataclass().
                    #
                    # embrace.mapobject:
                    # special case this because we can't statically infer the type
                    # of mapobject.passthrough
                    #
                    # In both cases typing.Any is the best we can do.
                    symbols.append(f"{base_aliases['typing']}.Any")
                elif mod == "builtins":
                    symbols.append(name)
                else:
                    if mod not in project_aliases:
                        project_aliases[mod] = f"_{next(c)}"
                    symbols.append(f"{project_aliases[mod]}.{name}")
        return symbols

    return symbol_aliases


class AttrDict(dict):
    def __getattr__(self, name):
        return self[name]


class ImportDict(dict):
    def __getattr__(self, name):
        return super().__getitem__(name.replace("_", "."))


def render_template(
    template: str,
    context: collections.abc.Mapping,
):
    tag_pattern = re.compile(
        r"""
        (
            {{ .*? }}
            | {% (?:[^%]|%[^}])* %}(?:[ \t]*\n)
            | {% (?:[^%]|%[^}])* %}
        )
        """,
        re.DOTALL | re.VERBOSE
    )
    template = dedent(template)
    tokens = [token for token in tag_pattern.split(template) if token]
    code = _parse_tokens(tokens)
    ns = dict(context)
    exec(code, ns)
    return "".join(ns["__output__"][0])


def _parse_tokens(
    tokens: collections.abc.Sequence[str],
    substitution_pattern=re.compile(r"{{\s*(.*?)\s*}}"),
    block_pattern=re.compile(
        r"{%\s(if |for |elif |else|capture |endcapture|end)(.*?)\s*%}\s*"
    ),
) -> str:
    captures = []
    lines = [
        "__output__ = [[]]\n"
        "def __out(s):\n"
        "    __output__[-1].append(str(s))\n"
        "\n",
    ]
    indent = 0

    def add_line(line, indent_delta=0, newline=True):
        nonlocal indent
        if lines[-1].endswith(":\n"):
            indent += 1
        indent += indent_delta
        if indent < 0:
            raise SyntaxError("Unexpected end tag")
        lines.append(f"{'    ' * indent}{line}\n")

    for token in tokens:
        indent_delta = 0
        if mo := substitution_pattern.match(token):
            add_line(f"__out({mo.group(1)})")

        elif mo := block_pattern.match(token):
            command, content = mo.groups()
            if command in {"elif ", "else", "end"}:
                indent_delta = -1
            if command == "end":
                indent -= 1
            elif command == "capture ":
                captures.append(content)
                add_line("__output__.append([])")
            elif command == "endcapture":
                capture = captures.pop()
                add_line(f"{capture} = ''.join(__output__.pop())")
                add_line(f"__out({capture})")
            else:
                add_line(f"{command}{content}:", indent_delta)
        else:
            add_line(f"__out({str(token)!r})")

    return "".join(lines)


def dedent(s: str):
    indent = None
    result = []
    for line in s.split("\n"):
        if indent is None:
            if line.strip():
                indent = re.match(r"^\s*", line).group(0)  # type: ignore
        if indent is None:
            result.append(line)
        else:
            result.append(line.removeprefix(indent))
    return "\n".join(result)


if __name__ == "__main__":
    import os
    import pathlib
    import argparse
    import tempfile

    argparser = argparse.ArgumentParser(
        description="Compile SQL queries into a single python module"
    )
    argparser.add_argument("--output", "-o", default="-", type=pathlib.Path)
    argparser.add_argument("dir", nargs="+")
    args = argparser.parse_args()
    m = Module()
    for d in args.dir:
        m.load_dir(d)
    s = generate_module_source(m)
    if args.output != "-":
        tmp = tempfile.NamedTemporaryFile(
            "w", encoding="utf-8", dir=args.output.parent, delete=False
        )
        tmp.write(s)
        tmp.close()
        os.rename(tmp.name, args.output)
    else:
        print(s)
