from promplate.prompt import *#

class CodeBuilder:
    """Build source code conveniently."""

    def __init__(self, indent_level=0, indent_str="\t"):
        self.code = []
        self.indent_level = indent_level
        self.indent_str = indent_str

    def __str__(self):
        return "".join(map(str, self.code))

    def add_line(self, line):
        """Add a line of source to the code."""
        self.code.extend((self.indent_str * self.indent_level, line, "\n"))

        return self

    def add_section(self):
        """Add a section, a sub-CodeBuilder."""
        section = CodeBuilder(self.indent_level, self.indent_str)
        self.code.append(section)

        return section

    def indent(self):
        """Increase the current indent for following lines."""
        self.indent_level += 1

        return self

    def dedent(self):
        """Decrease the current indent for following lines."""
        self.indent_level -= 1

        return self

    def get_render_function(self) -> FunctionType:
        """Execute the code, and return a dict of globals it defines."""
        assert self.indent_level == 0
        python_source = str(self)
        global_namespace = {}
        exec(python_source, global_namespace)
        return global_namespace["render"]

class Message(TypedDict):
    role: Role
    content: str
    name: NotRequired[str]

class MessageBuilder:
    _initializing = True
    __slots__ = ("role", "content", "name")

    def __init__(self, role: Role, /, content: str = "", name: str | None = None):
        self.role: Role = role
        self.content = content
        self.name = name

    def __repr__(self):
        if self.name is not None:
            return f"<| {self.role} {self.name} |>"
        return f"<| {self.role} |>"

    def __getitem__(self, key):
        return getattr(self, key)

    def __setitem__(self, key, value):
        return setattr(self, key, value)

    def __setattr__(self, key, value):
        if not self._initializing:
            assert self is not U and self is not A and self is not S
            assert isinstance(value, str)
        return super().__setattr__(key, value)

    def __matmul__(self, name: str):
        assert isinstance(name, str) and name
        return self.__class__(self.role, self.content, name)

    def dict(self) -> Message:
        if self.name:
            return {"role": self.role, "content": self.content, "name": self.name}
        return {"role": self.role, "content": self.content}

    def __gt__(self, content: str) -> Message:
        assert isinstance(content, str)
        if self.name:
            return {"role": self.role, "content": content, "name": self.name}
        return {"role": self.role, "content": content}

class CodeBuilder:
    """Build source code conveniently."""

    def __init__(self, indent_level=0, indent_str="\t"):
        self.code = []
        self.indent_level = indent_level
        self.indent_str = indent_str

    def __str__(self):
        return "".join(map(str, self.code))

    def add_line(self, line):
        """Add a line of source to the code."""
        self.code.extend((self.indent_str * self.indent_level, line, "\n"))

        return self

    def add_section(self):
        """Add a section, a sub-CodeBuilder."""
        section = CodeBuilder(self.indent_level, self.indent_str)
        self.code.append(section)

        return section

    def indent(self):
        """Increase the current indent for following lines."""
        self.indent_level += 1

        return self

    def dedent(self):
        """Decrease the current indent for following lines."""
        self.indent_level -= 1

        return self

    def get_render_function(self) -> FunctionType:
        """Execute the code, and return a dict of globals it defines."""
        assert self.indent_level == 0
        python_source = str(self)
        global_namespace = {}
        exec(python_source, global_namespace)
        return global_namespace["render"]

class AutoNaming:
    def __new__(cls, *args, **kwargs):
        obj = super().__new__(cls)
        obj._bind_frame()
        return obj

    def _bind_frame(self):
        from inspect import currentframe

        self._frame = currentframe()

    @cached_property
    def _name(self):
        f = self._frame
        if f and f.f_back and (frame := f.f_back.f_back):
            for name, var in frame.f_locals.items():
                if var is self:
                    return name

    @property
    def class_name(self):
        return self.__class__.__name__

    fallback_name = class_name

    @property
    def name(self):
        return self._name or self.fallback_name

    @name.setter
    def name(self, name):
        self._name = name
        self._frame = None

    @name.deleter
    def name(self):
        del self._name

    def __repr__(self):
        if self._name:
            return f"<{self.class_name} {self.name}>"
        else:
            return f"<{self.class_name}>"

    def __str__(self):
        return f"<{self.name}>"

class Component(Protocol):
    def render(self, context: Context) -> str: ...

    async def arender(self, context: Context) -> str: ...

class TemplateCore(AutoNaming):
    """A simple template compiler, for a jinja2-like syntax."""

    def __init__(self, text: str):
        """Construct a Templite with the given `text`."""

        self.text = text

    def _flush(self):
        for line in self._buffer:
            self._builder.add_line(line)
        self._buffer.clear()

    @staticmethod
    def _unwrap_token(token: str):
        return dedent(token.strip()[2:-2].strip("-")).strip()

    def _on_literal_token(self, token: str):
        self._buffer.append(f"__append__({repr(token)})")

    def _on_eval_token(self, token):
        token = self._unwrap_token(token)
        if "\n" in token:
            mod = parse(token)
            [*rest, last] = mod.body
            assert isinstance(last, Expr), "{{ }} block must end with an expression, or you should use {# #} block"
            self._buffer.extend(unparse(rest).splitlines())  # type: ignore
            exp = unparse(last)
        else:
            exp = token
        self._buffer.append(f"__append__({exp})")

    def _on_exec_token(self, token):
        self._buffer.extend(self._unwrap_token(token).splitlines())

    def _on_special_token(self, token, sync: bool):
        inner = self._unwrap_token(token)

        if inner.startswith("end"):
            last = self._ops_stack.pop()
            assert last == inner.removeprefix("end")
            self._flush()
            self._builder.dedent()

        else:
            op = inner.split(" ", 1)[0]

            if op == "if" or op == "for" or op == "while":
                self._ops_stack.append(op)
                self._flush()
                self._builder.add_line(f"{inner}:")
                self._builder.indent()

            elif op == "else" or op == "elif":
                self._flush()
                self._builder.dedent()
                self._builder.add_line(f"{inner}:")
                self._builder.indent()

            else:
                params: str = self._make_context(inner)
                if sync:
                    self._buffer.append(f"__append__({op}.render({params}))")
                else:
                    self._buffer.append(f"__append__(await {op}.arender({params}))")

    @staticmethod
    def _make_context(text: str):
        """generate context parameter if specified otherwise use locals() by default"""
        if version_info >= (3, 13):
            return f"globals() | locals() | dict({text[text.index(' ') + 1:]})" if " " in text else "globals() | locals()"
        return f"locals() | dict({text[text.index(' ') + 1:]})" if " " in text else "locals()"

    def compile(self, sync=True, indent_str="\t"):
        self._buffer = []
        self._ops_stack = []
        self._builder = get_base_builder(sync, indent_str)

        for token in split_template_tokens(self.text):
            if not token:
                continue
            s_token = token.strip()
            if s_token.startswith("{{") and s_token.endswith("}}"):
                self._on_eval_token(token)
            elif s_token.startswith("{#") and s_token.endswith("#}"):
                self._on_exec_token(token)
            elif s_token.startswith("{%") and s_token.endswith("%}") and "\n" not in s_token:
                self._on_special_token(token, sync)
            else:
                self._on_literal_token(token)

        if self._ops_stack:
            raise SyntaxError(self._ops_stack)

        self._flush()
        self._builder.add_line("return ''.join(map(str, __parts__))")
        self._builder.dedent()

    error_handling: Literal["linecache", "tempfile", "file"]
    if version_info >= (3, 13):
        error_handling = "linecache"
    else:
        error_handling = "tempfile" if __debug__ else "file"

    def _patch_for_error_handling(self, sync: bool):
        match self.error_handling:
            case "linecache":
                add_linecache(self.name, partial(self.get_script, sync, "\t"))
            case "file" | "tempfile":
                file = save_tempfile(self.name, self.get_script(sync, "\t"), self.error_handling == "tempfile")
                sys_path.append(str(file.parent))

    @cached_property
    def _render_code(self):
        self.compile()
        return self._builder.get_render_function().__code__.replace(co_filename=self.name, co_name="render")

    def render(self, context: Context) -> str:
        try:
            return eval(self._render_code, context)
        except Exception:
            self._patch_for_error_handling(sync=True)
            raise

    @cached_property
    def _arender_code(self):
        self.compile(sync=False)
        return self._builder.get_render_function().__code__.replace(co_filename=self.name, co_name="arender")

    async def arender(self, context: Context) -> str:
        try:
            return await eval(self._arender_code, context)
        except Exception:
            self._patch_for_error_handling(sync=False)
            raise

    def get_script(self, sync=True, indent_str="    "):
        """compile template string into python script"""
        self.compile(sync, indent_str)
        return str(self._builder)

class Loader(AutoNaming):
    @classmethod
    def read(cls, path: str | Path, encoding="utf-8"):
        path = Path(path)
        obj = cls(path.read_text(encoding))
        obj.name = path.stem
        return obj

    @classmethod
    async def aread(cls, path: str | Path, encoding="utf-8"):
        from aiofiles import open

        async with open(path, encoding=encoding) as f:
            content = await f.read()

        path = Path(path)
        obj = cls(content)
        obj.name = path.stem
        return obj

    @classmethod
    def _patch_kwargs(cls, kwargs: dict) -> dict:
        return {"headers": {"User-Agent": get_user_agent(cls)}} | kwargs

    @staticmethod
    def _join_url(url: str):
        if url.startswith("http"):
            return url

        from urllib.parse import urljoin

        return urljoin("https://promplate.dev/", url)

    @classmethod
    def fetch(cls, url: str, **kwargs):
        from .utils import _get_client

        response = _get_client().get(cls._join_url(url), **cls._patch_kwargs(kwargs))
        obj = cls(response.raise_for_status().text)
        obj.name = Path(url).stem
        return obj

    @classmethod
    async def afetch(cls, url: str, **kwargs):
        from .utils import _get_aclient

        response = await _get_aclient().get(cls._join_url(url), **cls._patch_kwargs(kwargs))
        obj = cls(response.raise_for_status().text)
        obj.name = Path(url).stem
        return obj

class SafeChainMapContext(ChainMap, dict):  # type: ignore
    pass

class Template(TemplateCore, Loader):
    def __init__(self, text: str, /, context: Context | None = None):
        super().__init__(text)
        self.context = {} if context is None else context

    def render(self, context: Context | None = None):
        if context is None:
            context = SafeChainMapContext({}, self.context)
        else:
            context = SafeChainMapContext({}, context, self.context)

        return super().render(context)

    async def arender(self, context: Context | None = None):
        if context is None:
            context = SafeChainMapContext({}, self.context)
        else:
            context = SafeChainMapContext({}, context, self.context)

        return await super().arender(context)