"""Function statement compilation for Kida compiler.

Provides mixin for compiling function statements (def, call_block, slot).

Uses inline TYPE_CHECKING declarations for host attributes.
See: plan/rfc-mixin-protocol-typing.md
"""

from __future__ import annotations

import ast
import logging
from collections.abc import Sequence
from typing import TYPE_CHECKING

from kida.nodes.output import Data

if TYPE_CHECKING:
    from kida.nodes import CallBlock, Def, Node, Region, Slot

logger = logging.getLogger(__name__)


def _slot_body_is_empty(slot_body: Sequence[object]) -> bool:
    """Treat whitespace-only Data nodes the same as an empty slot body."""
    return not slot_body or all(
        isinstance(child, Data) and not child.value.strip() for child in slot_body
    )


class FunctionCompilationMixin:
    """Mixin for compiling function statements.

    Host attributes and cross-mixin dependencies are declared via inline
    TYPE_CHECKING blocks.

    """

    # ─────────────────────────────────────────────────────────────────────────
    # Host attributes and cross-mixin dependencies (type-check only)
    # ─────────────────────────────────────────────────────────────────────────
    if TYPE_CHECKING:
        # Host attributes (from Compiler.__init__)
        _locals: set[str]
        _def_names: set[str]
        _def_caller_stack: list[ast.expr]

        # From ExpressionCompilationMixin
        def _compile_expr(self, node: Node, store: bool = False) -> ast.expr: ...

        # From Compiler core
        def _compile_node(self, node: Node) -> list[ast.stmt]: ...

    @staticmethod
    def _parse_annotation(raw: str) -> ast.expr | None:
        """Convert a raw annotation string to a Python AST expression node.

        Uses ``ast.parse`` in eval mode to parse annotation text like
        ``"str | None"`` into the corresponding AST. Falls back to None
        for malformed annotations.

        Args:
            raw: Annotation text from the template source.

        Returns:
            AST expression node, or None if parsing fails.
        """
        try:
            return ast.parse(raw, mode="eval").body
        except SyntaxError:
            logger.warning("Malformed type annotation in {%% def %%}: %r", raw)
            return None

    def _make_callable_preamble(self, *, include_scope_stack: bool = False) -> list[ast.stmt]:
        """Build common local runtime preamble for callable codegen paths."""
        stmts: list[ast.stmt] = [
            ast.Assign(
                targets=[ast.Name(id="_e", ctx=ast.Store())],
                value=ast.Name(id="_escape", ctx=ast.Load()),
            ),
            ast.Assign(
                targets=[ast.Name(id="_s", ctx=ast.Store())],
                value=ast.Name(id="_str", ctx=ast.Load()),
            ),
            ast.Assign(
                targets=[ast.Name(id="buf", ctx=ast.Store())],
                value=ast.List(elts=[], ctx=ast.Load()),
            ),
            ast.Assign(
                targets=[ast.Name(id="_append", ctx=ast.Store())],
                value=ast.Attribute(
                    value=ast.Name(id="buf", ctx=ast.Load()),
                    attr="append",
                    ctx=ast.Load(),
                ),
            ),
            ast.Assign(
                targets=[ast.Name(id="_acc", ctx=ast.Store())],
                value=ast.Call(
                    func=ast.Name(id="_get_accumulator", ctx=ast.Load()),
                    args=[],
                    keywords=[],
                ),
            ),
        ]
        if include_scope_stack:
            stmts.append(
                ast.Assign(
                    targets=[ast.Name(id="_scope_stack", ctx=ast.Store())],
                    value=ast.List(elts=[], ctx=ast.Load()),
                )
            )
        return stmts

    def _compile_def(self, node: Def) -> list[ast.stmt]:
        """Compile {% def name(params) %}...{% enddef %.

        Kida functions have true lexical scoping - they can access variables
        from their enclosing scope, unlike Jinja2 macros.

        Generates:
            def _def_name(arg1, arg2=default, *, _caller=None, _outer_ctx=ctx):
                buf = []
                ctx = {**_outer_ctx, 'arg1': arg1, 'arg2': arg2}
                if _caller:
                    ctx['caller'] = _caller
                ... body ...
                return Markup(''.join(buf))
            ctx['name'] = _def_name
        """
        def_name = node.name
        func_name = f"_def_{def_name}"
        # Track for profiling: FuncCall to this name will be instrumented
        self._def_names.add(def_name)

        # Build function arguments with optional type annotations
        args_list = [
            ast.arg(
                arg=p.name,
                annotation=(self._parse_annotation(p.annotation) if p.annotation else None),
            )
            for p in node.params
        ]
        defaults = [self._compile_expr(d) for d in node.defaults]

        # Build vararg and kwarg AST nodes
        vararg_node = ast.arg(arg=node.vararg) if node.vararg else None
        kwarg_node = ast.arg(arg=node.kwarg) if node.kwarg else None

        # Build context dict entries: regular args + vararg + kwarg
        param_names = [p.name for p in node.params]
        ctx_keys: list[ast.expr | None] = [ast.Constant(value=name) for name in param_names]
        ctx_values: list[ast.expr] = [ast.Name(id=name, ctx=ast.Load()) for name in param_names]

        if node.vararg:
            ctx_keys.append(ast.Constant(value=node.vararg))
            ctx_values.append(ast.Name(id=node.vararg, ctx=ast.Load()))

        if node.kwarg:
            ctx_keys.append(ast.Constant(value=node.kwarg))
            ctx_values.append(ast.Name(id=node.kwarg, ctx=ast.Load()))

        # Build function body
        func_body: list[ast.stmt] = [
            *self._make_callable_preamble(),
            # Create local context: ctx = {**_outer_ctx, 'arg1': arg1, ...}
            ast.Assign(
                targets=[ast.Name(id="ctx", ctx=ast.Store())],
                value=ast.Dict(
                    keys=[None, None],  # Spread operators
                    values=[
                        ast.Name(id="_outer_ctx", ctx=ast.Load()),
                        ast.Dict(keys=ctx_keys, values=ctx_values),
                    ],
                ),
            ),
            # if _caller: ctx['caller'] = _caller
            ast.If(
                test=ast.Name(id="_caller", ctx=ast.Load()),
                body=[
                    ast.Assign(
                        targets=[
                            ast.Subscript(
                                value=ast.Name(id="ctx", ctx=ast.Load()),
                                slice=ast.Constant(value="caller"),
                                ctx=ast.Store(),
                            )
                        ],
                        value=ast.Name(id="_caller", ctx=ast.Load()),
                    )
                ],
                orelse=[],
            ),
            # ctx['has_slot'] = lambda: _caller is not None
            # Allows {% if has_slot() %} guards inside def bodies.
            ast.Assign(
                targets=[
                    ast.Subscript(
                        value=ast.Name(id="ctx", ctx=ast.Load()),
                        slice=ast.Constant(value="has_slot"),
                        ctx=ast.Store(),
                    )
                ],
                value=ast.Lambda(
                    args=ast.arguments(
                        posonlyargs=[],
                        args=[],
                        vararg=None,
                        kwonlyargs=[],
                        kw_defaults=[],
                        kwarg=None,
                        defaults=[],
                    ),
                    body=ast.Compare(
                        left=ast.Name(id="_caller", ctx=ast.Load()),
                        ops=[ast.IsNot()],
                        comparators=[ast.Constant(value=None)],
                    ),
                ),
            ),
        ]

        # Add args to locals for direct access
        for p in node.params:
            self._locals.add(p.name)
        if node.vararg:
            self._locals.add(node.vararg)
        if node.kwarg:
            self._locals.add(node.kwarg)

        # Compile function body
        # Macros (def) are always sync functions — disable async mode
        # to prevent async for/await inside the def body.
        # Lexical caller scoping: capture _caller to _def_caller so slot functions
        # can reference it without shadowing by the inner _caller wrapper.
        saved_async = getattr(self, "_async_mode", False)
        if saved_async:
            self._async_mode = False
        func_body.append(
            ast.Assign(
                targets=[ast.Name(id="_def_caller", ctx=ast.Store())],
                value=ast.Name(id="_caller", ctx=ast.Load()),
            )
        )
        self._def_caller_stack.append(ast.Name(id="_def_caller", ctx=ast.Load()))
        try:
            for child in node.body:
                func_body.extend(self._compile_node(child))
        finally:
            self._def_caller_stack.pop()
        if saved_async:
            self._async_mode = saved_async

        # Remove args from locals
        for p in node.params:
            self._locals.discard(p.name)
        if node.vararg:
            self._locals.discard(node.vararg)
        if node.kwarg:
            self._locals.discard(node.kwarg)

        # return _Markup(''.join(buf))
        func_body.append(
            ast.Return(
                value=ast.Call(
                    func=ast.Name(id="_Markup", ctx=ast.Load()),
                    args=[
                        ast.Call(
                            func=ast.Attribute(
                                value=ast.Constant(value=""),
                                attr="join",
                                ctx=ast.Load(),
                            ),
                            args=[ast.Name(id="buf", ctx=ast.Load())],
                            keywords=[],
                        ),
                    ],
                    keywords=[],
                ),
            )
        )

        # Create function with _caller and _outer_ctx as keyword-only args
        # When **kwargs is used, _caller and _outer_ctx must use a different
        # mechanism since Python only allows one **kwarg. We place them as
        # keyword-only args that come before the **kwarg isn't possible in
        # Python syntax. Instead, _caller and _outer_ctx are always
        # keyword-only args, and user **kwargs is separate.
        #
        # Strategy: When **kwargs is present, we use a wrapper that splits
        # the internal kwargs (_caller, _outer_ctx) from user kwargs.
        # When no **kwargs, we use the simple approach.
        kwonlyargs = [
            ast.arg(arg="_caller"),
            ast.arg(arg="_outer_ctx"),
        ]
        kw_defaults: list[ast.expr | None] = [
            ast.Constant(value=None),  # _caller=None
            ast.Name(id="ctx", ctx=ast.Load()),  # _outer_ctx=ctx
        ]

        func_def = ast.FunctionDef(
            name=func_name,
            args=ast.arguments(
                posonlyargs=[],
                args=args_list,
                vararg=vararg_node,
                kwonlyargs=kwonlyargs,
                kw_defaults=kw_defaults,
                kwarg=kwarg_node,
                defaults=defaults,
            ),
            body=func_body,
            decorator_list=[],
            returns=None,
        )

        # Assign to context: ctx['name'] = _def_name
        assign = ast.Assign(
            targets=[
                ast.Subscript(
                    value=ast.Name(id="ctx", ctx=ast.Load()),
                    slice=ast.Constant(value=def_name),
                    ctx=ast.Store(),
                )
            ],
            value=ast.Name(id=func_name, ctx=ast.Load()),
        )

        return [func_def, assign]

    def _compile_call_block(self, node: CallBlock) -> list[ast.stmt]:
        """Compile {% call func(args) %}...{% endcall %.

        Builds a caller that supports named slots. Passes _caller(slot="default")
        so both _caller() and _caller("header_actions") work.
        """
        stmts: list[ast.stmt] = []
        slots = node.slots

        # Build one function per slot
        slot_fns: dict[str, str] = {}
        for slot_name, slot_body in slots.items():
            safe_name = slot_name.replace("-", "_")
            fn_name = f"_caller_{safe_name}"

            caller_body: list[ast.stmt] = self._make_callable_preamble()

            saved_async_cb = getattr(self, "_async_mode", False)
            if saved_async_cb:
                self._async_mode = False
            outer = self._def_caller_stack[-1] if self._def_caller_stack else None
            # Use _outer_caller param to avoid shadowing by inner _caller wrapper
            if outer is not None:
                self._outer_caller_expr = ast.Name(id="_outer_caller", ctx=ast.Load())
                # Delegate to outer caller when slot body is empty (fixes nested macro slot passthrough)
                if _slot_body_is_empty(slot_body):
                    caller_body.append(
                        ast.If(
                            test=ast.Compare(
                                left=ast.Name(id="_outer_caller", ctx=ast.Load()),
                                ops=[ast.IsNot()],
                                comparators=[ast.Constant(value=None)],
                            ),
                            body=[
                                ast.Return(
                                    value=ast.Call(
                                        func=ast.Name(id="_outer_caller", ctx=ast.Load()),
                                        args=[ast.Constant(value=slot_name)],
                                        keywords=[],
                                    )
                                )
                            ],
                            orelse=[],
                        )
                    )
            try:
                for child in slot_body:
                    caller_body.extend(self._compile_node(child))
            finally:
                if outer is not None:
                    self._outer_caller_expr = None
            if saved_async_cb:
                self._async_mode = saved_async_cb

            caller_body.append(
                ast.Return(
                    value=ast.Call(
                        func=ast.Name(id="_Markup", ctx=ast.Load()),
                        args=[
                            ast.Call(
                                func=ast.Attribute(
                                    value=ast.Constant(value=""),
                                    attr="join",
                                    ctx=ast.Load(),
                                ),
                                args=[ast.Name(id="buf", ctx=ast.Load())],
                                keywords=[],
                            ),
                        ],
                        keywords=[],
                    ),
                )
            )

            slot_fns[slot_name] = fn_name
            # _scope_stack required for _lookup_scope; _outer_caller for def nesting
            slot_args: list[ast.arg] = [ast.arg(arg="_scope_stack")]
            slot_defaults: list[ast.expr] = []
            if outer is not None:
                slot_args.append(ast.arg(arg="_outer_caller"))
                slot_defaults.append(outer)
            stmts.append(
                ast.FunctionDef(
                    name=fn_name,
                    args=ast.arguments(
                        posonlyargs=[],
                        args=slot_args,
                        vararg=None,
                        kwonlyargs=[],
                        kw_defaults=[],
                        kwarg=None,
                        defaults=slot_defaults,
                    ),
                    body=caller_body,
                    decorator_list=[],
                    returns=None,
                )
            )

        # Build _caller(slot="default") wrapper
        # def _caller(slot="default"):
        #     f = _caller_slots.get(slot)
        #     return f(_scope_stack) if f else _Markup("")
        wrapper_body = [
            ast.Assign(
                targets=[ast.Name(id="_f", ctx=ast.Store())],
                value=ast.Call(
                    func=ast.Attribute(
                        value=ast.Name(id="_caller_slots", ctx=ast.Load()),
                        attr="get",
                        ctx=ast.Load(),
                    ),
                    args=[ast.Name(id="slot", ctx=ast.Load())],
                    keywords=[],
                ),
            ),
            ast.Return(
                value=ast.IfExp(
                    test=ast.Name(id="_f", ctx=ast.Load()),
                    body=ast.Call(
                        func=ast.Name(id="_f", ctx=ast.Load()),
                        args=[ast.Name(id="_scope_stack", ctx=ast.Load())],
                        keywords=[],
                    ),
                    orelse=ast.Call(
                        func=ast.Name(id="_Markup", ctx=ast.Load()),
                        args=[ast.Constant(value="")],
                        keywords=[],
                    ),
                ),
            ),
        ]
        stmts.append(
            ast.Assign(
                targets=[ast.Name(id="_caller_slots", ctx=ast.Store())],
                value=ast.Dict(
                    keys=[ast.Constant(value=k) for k in slot_fns],
                    values=[ast.Name(id=v, ctx=ast.Load()) for v in slot_fns.values()],
                ),
            )
        )
        # Use _caller_wrapper to avoid shadowing the def's _caller parameter
        # (which would cause UnboundLocalError when _def_caller = _caller runs)
        # Wrapper takes (slot, _scope_stack) so slot functions get scope for _lookup_scope
        stmts.append(
            ast.FunctionDef(
                name="_caller_wrapper",
                args=ast.arguments(
                    posonlyargs=[],
                    args=[ast.arg(arg="slot"), ast.arg(arg="_scope_stack")],
                    vararg=None,
                    kwonlyargs=[],
                    kw_defaults=[],
                    kwarg=None,
                    defaults=[ast.Constant(value="default")],
                ),
                body=wrapper_body,
                decorator_list=[],
                returns=None,
            )
        )
        # Lambda captures _scope_stack from render scope; macro calls _caller() or _caller(slot)
        stmts.append(
            ast.Assign(
                targets=[ast.Name(id="_caller_with_scope", ctx=ast.Store())],
                value=ast.Lambda(
                    args=ast.arguments(
                        posonlyargs=[],
                        args=[ast.arg(arg="slot")],
                        vararg=None,
                        kwonlyargs=[],
                        kw_defaults=[],
                        kwarg=None,
                        defaults=[ast.Constant(value="default")],
                    ),
                    body=ast.Call(
                        func=ast.Name(id="_caller_wrapper", ctx=ast.Load()),
                        args=[
                            ast.Name(id="slot", ctx=ast.Load()),
                            ast.Name(id="_scope_stack", ctx=ast.Load()),
                        ],
                        keywords=[],
                    ),
                ),
            )
        )

        # Compile the call expression and add _caller keyword argument
        # Suppress macro instrumentation so we get a raw ast.Call back
        saved_skip = getattr(self, "_skip_macro_instrumentation", False)
        self._skip_macro_instrumentation = True
        call_expr = self._compile_expr(node.call)
        self._skip_macro_instrumentation = saved_skip

        # If it's a function call, add _caller keyword (lambda binds _scope_stack)
        if isinstance(call_expr, ast.Call):
            call_expr.keywords.append(
                ast.keyword(arg="_caller", value=ast.Name(id="_caller_with_scope", ctx=ast.Load()))
            )
        else:
            # Wrap in a call with _caller
            call_expr = ast.Call(
                func=call_expr,
                args=[],
                keywords=[
                    ast.keyword(
                        arg="_caller",
                        value=ast.Name(id="_caller_with_scope", ctx=ast.Load()),
                    )
                ],
            )

        # _append(result)
        stmts.append(
            ast.Expr(
                value=ast.Call(
                    func=ast.Name(id="_append", ctx=ast.Load()),
                    args=[call_expr],
                    keywords=[],
                ),
            )
        )

        return stmts

    def _make_region_function(self, name: str, node: Region) -> ast.FunctionDef:
        """Generate module-level region callable: _region_name(params, *, _outer_ctx) -> Markup.

        Regions compile to both a callable (this) and a block wrapper.
        No _caller/slots — regions are parameterized renderable units only.
        """
        func_name = f"_region_{name}"
        param_names = [p.name for p in node.params]
        ctx_keys: list[ast.expr | None] = [ast.Constant(value=n) for n in param_names]
        ctx_values = [ast.Name(id=n, ctx=ast.Load()) for n in param_names]
        if node.vararg:
            ctx_keys.append(ast.Constant(value=node.vararg))
            ctx_values.append(ast.Name(id=node.vararg, ctx=ast.Load()))
        if node.kwarg:
            ctx_keys.append(ast.Constant(value=node.kwarg))
            ctx_values.append(ast.Name(id=node.kwarg, ctx=ast.Load()))

        args_list = [
            ast.arg(
                arg=p.name,
                annotation=(self._parse_annotation(p.annotation) if p.annotation else None),
            )
            for p in node.params
        ]
        defaults = [self._compile_expr(d) for d in node.defaults]
        vararg_node = ast.arg(arg=node.vararg) if node.vararg else None
        kwarg_node = ast.arg(arg=node.kwarg) if node.kwarg else None

        func_body: list[ast.stmt] = [
            *self._make_callable_preamble(include_scope_stack=True),
            ast.Assign(
                targets=[ast.Name(id="ctx", ctx=ast.Store())],
                value=ast.Dict(
                    keys=[None, None],
                    values=[
                        ast.Name(id="_outer_ctx", ctx=ast.Load()),
                        ast.Dict(keys=ctx_keys, values=ctx_values),
                    ],
                ),
            ),
        ]

        for p in node.params:
            self._locals.add(p.name)
        if node.vararg:
            self._locals.add(node.vararg)
        if node.kwarg:
            self._locals.add(node.kwarg)

        saved_async = getattr(self, "_async_mode", False)
        if saved_async:
            self._async_mode = False
        try:
            for child in node.body:
                func_body.extend(self._compile_node(child))
        finally:
            for p in node.params:
                self._locals.discard(p.name)
            if node.vararg:
                self._locals.discard(node.vararg)
            if node.kwarg:
                self._locals.discard(node.kwarg)
            if saved_async:
                self._async_mode = True

        func_body.append(
            ast.Return(
                value=ast.Call(
                    func=ast.Name(id="_Markup", ctx=ast.Load()),
                    args=[
                        ast.Call(
                            func=ast.Attribute(
                                value=ast.Constant(value=""),
                                attr="join",
                                ctx=ast.Load(),
                            ),
                            args=[ast.Name(id="buf", ctx=ast.Load())],
                            keywords=[],
                        ),
                    ],
                    keywords=[],
                ),
            )
        )

        return ast.FunctionDef(
            name=func_name,
            args=ast.arguments(
                posonlyargs=[],
                args=args_list,
                vararg=vararg_node,
                kwonlyargs=[ast.arg(arg="_outer_ctx")],
                kw_defaults=[None],  # Required kw-only (no default)
                kwarg=kwarg_node,
                defaults=defaults,
            ),
            body=func_body,
            decorator_list=[],
            returns=None,
        )

    def _compile_region(self, node: Region) -> list[ast.stmt]:
        """Compile {% region name(params) %} for template body — emit ctx assign only."""
        func_name = f"_region_{node.name}"
        return [
            ast.Assign(
                targets=[
                    ast.Subscript(
                        value=ast.Name(id="ctx", ctx=ast.Load()),
                        slice=ast.Constant(value=node.name),
                        ctx=ast.Store(),
                    )
                ],
                value=ast.Name(id=func_name, ctx=ast.Load()),
            ),
        ]

    def _compile_slot(self, node: Slot) -> list[ast.stmt]:
        """Compile {% slot %} or {% slot name %.

        Renders the caller content for the given slot name.
        _caller(slot="default") supports both _caller() and _caller("name").
        """
        slot_name = node.name
        return [
            ast.If(
                test=ast.Call(
                    func=ast.Attribute(
                        value=ast.Name(id="ctx", ctx=ast.Load()),
                        attr="get",
                        ctx=ast.Load(),
                    ),
                    args=[ast.Constant(value="caller")],
                    keywords=[],
                ),
                body=[
                    ast.Expr(
                        value=ast.Call(
                            func=ast.Name(id="_append", ctx=ast.Load()),
                            args=[
                                ast.Call(
                                    func=ast.Subscript(
                                        value=ast.Name(id="ctx", ctx=ast.Load()),
                                        slice=ast.Constant(value="caller"),
                                        ctx=ast.Load(),
                                    ),
                                    args=[ast.Constant(value=slot_name)],
                                    keywords=[],
                                ),
                            ],
                            keywords=[],
                        ),
                    )
                ],
                orelse=[],
            )
        ]
