# -*- coding: utf-8 -*-
"""Monadic do-notation as a block macro.

Syntax::

    with monadic_do[M] as result:
        [x := mx,
         y := my(x),
         M.guard(...),
         M.unit(x + y)]

The body is a single list literal. Each item corresponds to one line of
a Haskell do-block. The **last item** is the final monadic expression
(any expression of type ``M a``, matching Haskell's last-line-of-do).
All **earlier items** are binds:

- ``name := mexpr`` — monadic bind: the unwrapped value is bound to
  ``name`` for subsequent lines.
- ``name << mexpr`` — legacy alternative for ``:=`` (same shapes
  ``letdoutil`` recognizes for ``let[]``).
- a bare ``mexpr`` — sequencing-only (Haskell's ``do { mx; ... }``): the
  result is threaded but discarded. The short-circuit behavior of the
  monad still applies (``Maybe(nil)``, ``Left``, empty ``List`` all
  cancel the rest of the chain).

Expands to a nested lambda-bind chain::

    result = mx >> (lambda x: my(x) >> (lambda _: M.guard(...) >> (lambda _: M.unit(x + y))))

**Placement in the xmas tree**: always the innermost ``with``. Its body
shape (a single list-literal statement) forbids lexically wrapping other
``with`` blocks inside it, and outer two-pass macros (``lazify``,
``continuations``, ``tco``, ``autocurry``, etc.) expand inner macros
between their two passes, which means they will correctly see and edit
the expanded bind chain.
"""

__all__ = ["monadic_do"]

from ast import List, Name, NamedExpr, BinOp, LShift, Expr, Assign, Store, arg, expr

from mcpyrate.quotes import macros, q, a, n  # noqa: F401

from mcpyrate import parametricmacro

from ..dynassign import dyn

from .letdoutil import canonize_bindings


@parametricmacro
def monadic_do(tree, *, args, syntax, expander, **kw):
    """[syntax, block] Monadic do-notation.

    See module docstring for usage, placement, and expansion.
    """
    if syntax != "block":
        raise SyntaxError("monadic_do is a block macro only")  # pragma: no cover

    # Require exactly one macro argument: the monad type.
    if len(args) != 1:
        raise SyntaxError(
            f"monadic_do expects exactly one macro argument (the monad type), got {len(args)}"
        )  # pragma: no cover

    # Require the `as` binding — this is where the result lands.
    result_var = kw.get("optional_vars", None)
    if result_var is None:
        raise SyntaxError(
            "monadic_do requires an as-binding: `with monadic_do[M] as result:`"
        )  # pragma: no cover
    if type(result_var) is not Name:
        raise SyntaxError(
            "monadic_do's as-binding must be a single name"
        )  # pragma: no cover

    with dyn.let(_macro_expander=expander):
        return _monadic_do(block_body=tree, monad_type=args[0], result_name=result_var.id)


def _monadic_do(block_body: list, monad_type: expr, result_name: str) -> list:
    # Expand inner macros first (outside-in), just like `forall` and `autoref` do.
    block_body = dyn._macro_expander.visit_recursively(block_body)

    # Body must be exactly one statement, an Expr wrapping a List literal.
    if len(block_body) != 1:
        raise SyntaxError(
            f"monadic_do body must be a single list-literal statement, got {len(block_body)} statements"
        )  # pragma: no cover
    stmt = block_body[0]
    if type(stmt) is not Expr or type(stmt.value) is not List:
        raise SyntaxError(
            "monadic_do body must be a single list literal `[bind, ..., final_expr]`"
        )  # pragma: no cover

    items = stmt.value.elts
    if not items:
        raise SyntaxError(
            "monadic_do body list must have at least one item (the final monadic expression)"
        )  # pragma: no cover

    # Split: all but the last are binds; the last is the final monadic expression.
    *binding_items, final_expr = items

    # Normalize bare expressions in the binds as synthetic `_ := expr` so they
    # look like sequencing-only bindings to `canonize_bindings`. Matches Haskell's
    # do-notation where a bare expression line is sequence-only (>>, not >>=).
    normalized = [
        item if _is_binding_form(item) else NamedExpr(target=Name(id="_", ctx=Store()), value=item)
        for item in binding_items
    ]

    # Parse via letdoutil — accepts := and <<.
    if normalized:
        canonical = canonize_bindings(normalized)  # [Tuple(elts=[Name(k), v]), ...]
        pairs = [(t.elts[0].id, t.elts[1]) for t in canonical]
    else:
        pairs = []

    # Build the bind chain, innermost-first:
    #   final_expr
    #   mz >> (lambda z: final_expr)
    #   my >> (lambda y: mz >> (lambda z: final_expr))
    #   mx >> (lambda x: my >> (lambda y: mz >> (lambda z: final_expr)))
    body = final_expr
    for name, mexpr in reversed(pairs):
        # lambda <name>: <body>
        lam = q[lambda: a[body]]
        lam.args.args = [arg(arg=name)]
        # <mexpr> >> <lam>
        body = q[a[mexpr] >> a[lam]]

    # Final assignment: `<result_name> = <body>`. This is a statement; we replace
    # the entire `with` body with it.
    assignment = Assign(targets=[Name(id=result_name, ctx=Store())], value=body)
    return [assignment]


def _is_binding_form(item) -> bool:
    """Return True if *item* is ``name := expr`` or ``name << expr`` (a let-style binding)."""
    if type(item) is NamedExpr and type(item.target) is Name:
        return True
    if type(item) is BinOp and type(item.op) is LShift and type(item.left) is Name:  # noqa: SIM103 -- keep cases visually separate
        return True
    return False