"""
Nexa Language — Auto-Formatter (v22)
`nexa fmt file.nexa` — canonical formatting, similar to gofmt.

Rules:
  - 4-space indentation
  - Paren-free control flow
  - Single space around binary operators
  - Trailing newline
  - One blank line between top-level definitions
"""
from __future__ import annotations
from typing import Any, List
from nexa_lang.ast_nodes import *


class Formatter:
    def __init__(self, indent=4):
        self._indent = indent
        self._level = 0
        self._out: List[str] = []
        self._prev_was_def = False

    def format(self, program: Program) -> str:
        for i, stmt in enumerate(program.statements):
            if i > 0 and isinstance(stmt, (FunctionDef, ClassDef)):
                self._out.append("")  # blank line before top-level defs
            self._stmt(stmt)
        return "\n".join(self._out) + "\n"

    def _ind(self) -> str:
        return " " * (self._indent * self._level)

    def _line(self, s: str):
        self._out.append(self._ind() + s)

    def _block(self, stmts: list):
        self._line("{")
        self._level += 1
        for stmt in stmts:
            self._stmt(stmt)
        self._level -= 1
        self._line("}")

    def _stmt(self, node):
        t = type(node).__name__
        m = getattr(self, f"_f_{t}", self._f_fallback)
        m(node)

    def _f_fallback(self, node):
        self._line(self._expr(node))

    def _expr(self, node) -> str:
        t = type(node).__name__
        m = getattr(self, f"_e_{t}", lambda n: repr(n))
        return m(node)

    # ── Statement formatters ──────────────────────────────────────────────────
    def _f_LetStatement(self, node):
        type_hint = f": {node.type_hint}" if getattr(node, 'type_hint', None) else ""
        self._line(f"let {node.name}{type_hint} = {self._expr(node.value)}")

    def _f_Assignment(self, node):
        self._line(f"{node.name} = {self._expr(node.value)}")

    def _f_ExprStatement(self, node):
        self._line(self._expr(node.expr))

    def _f_ReturnStatement(self, node):
        val = f" {self._expr(node.value)}" if node.value else ""
        self._line(f"return{val}")

    def _e_IfExpr(self, node) -> str:
        res = ["if " + self._expr(node.condition) + " {"]
        self._level += 1
        for s in node.then_block: 
            if type(s).__name__ == "ExprStatement": res.append(self._ind() + self._expr(s.expr))
            else: res.append(self._ind() + "<stmt>")
        self._level -= 1
        if node.else_block:
            if len(node.else_block) == 1 and type(node.else_block[0]).__name__ == "IfExpr":
                inner = self._e_IfExpr(node.else_block[0]).replace("\n", "\n" + self._ind())
                res.append("} else " + inner)
                return "\n".join(res)
            res.append(self._ind() + "} else {")
            self._level += 1
            for s in node.else_block: 
                if type(s).__name__ == "ExprStatement": res.append(self._ind() + self._expr(s.expr))
                else: res.append(self._ind() + "<stmt>")
            self._level -= 1
        res.append(self._ind() + "}")
        return "\n".join(res)

    def _f_ForStatement(self, node):
        self._line(f"for {node.var} in {self._expr(node.iterable)} {{")
        self._level += 1
        for s in node.body: self._stmt(s)
        self._level -= 1
        self._line("}")

    def _f_WhileStatement(self, node):
        self._line(f"while {self._expr(node.condition)} {{")
        self._level += 1
        for s in node.body: self._stmt(s)
        self._level -= 1
        self._line("}")

    def _f_LoopStatement(self, node):
        self._line("loop {")
        self._level += 1
        for s in node.body: self._stmt(s)
        self._level -= 1
        self._line("}")

    def _f_BreakStatement(self, node):    self._line("break")
    def _f_ContinueStatement(self, node): self._line("continue")

    def _f_FunctionDef(self, node):
        vis = (getattr(node, 'visibility', None) or '')
        vis_prefix = f"{vis} " if vis else ""
        tps = f"<{', '.join(getattr(node, 'type_params', []) or [])}>" if getattr(node, 'type_params', None) else ""
        params = ", ".join(node.params)
        ret = f" -> {node.return_type}" if getattr(node, 'return_type', None) else ""
        prefix = "async " if getattr(node, 'is_async', False) else ""
        self._line(f"{vis_prefix}{prefix}def {node.name}{tps}({params}){ret} {{")
        self._level += 1
        for s in node.body: self._stmt(s)
        self._level -= 1
        self._line("}")

    def _f_ClassDef(self, node):
        base = f"({node.base})" if node.base else ""
        impls = f" implements {', '.join(node.implements)}" if getattr(node, 'implements', None) else ""
        abstract = "abstract " if getattr(node, 'is_abstract', False) else ""
        tps = f"<{', '.join(getattr(node, 'type_params', []) or [])}>" if getattr(node, 'type_params', None) else ""
        self._line(f"{abstract}class {node.name}{tps}{base}{impls} {{")
        self._level += 1
        for s in node.body: self._stmt(s)
        self._level -= 1
        self._line("}")

    def _f_InterfaceDef(self, node):
        self._line(f"interface {node.name} {{")
        self._level += 1
        for m in node.method_names: self._line(f"def {m}()")
        self._level -= 1
        self._line("}")

    def _f_AbstractMethodDef(self, node):
        params = ", ".join(node.params)
        self._line(f"abstract def {node.name}({params})")

    def _f_ExportStatement(self, node):
        self._line(f"export {{ {', '.join(node.names)} }}")

    def _f_ImportStatement(self, node):
        alias = f" as {node.alias}" if node.alias else ""
        self._line(f"import {node.module}{alias}")

    def _f_ThrowStatement(self, node):
        self._line(f"throw {self._expr(node.value)}")

    def _e_TryExpr(self, node):
        res = ["try {"]
        self._level += 1
        for s in node.try_block: res.append(self._ind() + "<stmt>")
        self._level -= 1
        if node.catch_block:
            var = f"({node.catch_bind})" if getattr(node, "catch_bind", None) else "(e)"
            res.append(self._ind() + f"}} catch {var} {{")
            self._level += 1
            for s in node.catch_block: res.append(self._ind() + "<stmt>")
            self._level -= 1
        if node.finally_block:
            res.append(self._ind() + "} finally {")
            self._level += 1
            for s in node.finally_block: res.append(self._ind() + "<stmt>")
            self._level -= 1
        res.append(self._ind() + "}")
        return "\n".join(res)

    def _f_LoopStatement(self, node):
        self._line("loop {")
        self._level += 1
        for s in node.body: self._stmt(s)
        self._level -= 1
        self._line("}")

    # ── Expression formatters ──────────────────────────────────────────────────
    def _e_NumberLiteral(self, n):  return str(n.value)
    def _e_StringLiteral(self, n):  return json_str(n.value)
    def _e_BoolLiteral(self, n):    return "true" if n.value else "false"
    def _e_NullLiteral(self, n):    return "null"
    def _e_SelfRef(self, n):        return "self"
    def _e_Identifier(self, n):     return n.name
    def _e_NonNullAssertion(self, n): return self._expr(n.expr) + "!"

    def _e_BinaryOp(self, n):
        return f"({self._expr(n.left)} {n.op} {self._expr(n.right)})"

    def _e_UnaryOp(self, n):
        return f"({n.op}{self._expr(n.operand)})"

    def _e_FunctionCall(self, n):
        args = ", ".join(self._expr(a) for a in n.args)
        return f"{self._expr(n.callee)}({args})"

    def _e_MemberAccess(self, n):
        return f"{self._expr(n.obj)}.{n.member}"

    def _e_RangeExpr(self, n):
        op = "..=" if n.inclusive else ".."
        return f"{self._expr(n.start)}{op}{self._expr(n.end)}"

    def _e_ListLiteral(self, n):
        return "[" + ", ".join(self._expr(e) for e in n.elements) + "]"

    def _e_DictLiteral(self, n):
        pairs = ", ".join(f"{self._expr(k)}: {self._expr(v)}" for k, v in n.pairs)
        return "{" + pairs + "}"

    def _e_TernaryOp(self, n):
        return f"({self._expr(n.true_expr)} if {self._expr(n.condition)} else {self._expr(n.false_expr)})"

    def _e_MatchExpr(self, node):
        return "<match_expr>"

    def _e_BlockExpr(self, node):
        return "<block_expr>"


def json_str(s: str) -> str:
    import json
    return json.dumps(s)
