from __future__ import annotations

from .lexer import Lexer
from .parser import (
    Assignment,
    BinaryExpression,
    BreakStatement,
    ContinueStatement,
    FormattedString,
    FunctionCall,
    FunctionDeclaration,
    Identifier,
    IfStatement,
    IndexAssignment,
    IndexExpression,
    ListLiteral,
    Literal,
    ObjectLiteral,
    Program,
    ReturnStatement,
    ThrowStatement,
    TryStatement,
    UnaryExpression,
    WhileStatement,
    ForLoop,
    Parser,
)


_PREC = {
    "||": 1,
    "&&": 2,
    "==": 3,
    "!=": 3,
    ">": 4,
    "<": 4,
    ">=": 4,
    "<=": 4,
    "+": 5,
    "-": 5,
    "*": 6,
    "/": 6,
    "%": 6,
}


def _escape_string(value: str) -> str:
    value = value.replace("\\", "\\\\")
    value = value.replace('"', '\\"')
    value = value.replace("\n", "\\n")
    value = value.replace("\t", "\\t")
    return f'"{value}"'


def _needs_parens(child, parent_op: str, side: str) -> bool:
    if not isinstance(child, BinaryExpression):
        return False
    parent_prec = _PREC.get(parent_op, 0)
    child_prec = _PREC.get(child.op, 0)
    if child_prec < parent_prec:
        return True
    if child_prec > parent_prec:
        return False

    # Same precedence: be conservative on the right side for non-associative ops.
    if side == "right" and parent_op in ("-", "/", "%", "==", "!=", ">", "<", ">=", "<="):
        return True
    return False


def format_expr(expr) -> str:
    if isinstance(expr, Literal):
        if expr.value is None:
            return "kahi"
        if expr.value is True:
            return "ho"
        if expr.value is False:
            return "nahi"
        if isinstance(expr.value, str):
            return _escape_string(expr.value)
        return str(expr.value)

    if isinstance(expr, Identifier):
        return expr.name

    if isinstance(expr, ListLiteral):
        inner = ", ".join(format_expr(e) for e in expr.elements)
        return f"[{inner}]"

    if isinstance(expr, ObjectLiteral):
        parts = []
        for k, v in expr.entries:
            parts.append(f"{_escape_string(k)}: {format_expr(v)}")
        inner = ", ".join(parts)
        return "{ " + inner + " }" if parts else "{ }"

    if isinstance(expr, IndexExpression):
        return f"{format_expr(expr.target)}[{format_expr(expr.index)}]"

    if isinstance(expr, UnaryExpression):
        op = "navhe" if expr.op == "!" else expr.op
        # Our parser stores keyword-not as "!" too; prefer symbol form for formatting.
        if op == "navhe":
            op = "!"
        right = format_expr(expr.expr)
        return f"{op}{right}"

    if isinstance(expr, BinaryExpression):
        left_s = format_expr(expr.left)
        right_s = format_expr(expr.right)
        if _needs_parens(expr.left, expr.op, "left"):
            left_s = f"({left_s})"
        if _needs_parens(expr.right, expr.op, "right"):
            right_s = f"({right_s})"
        return f"{left_s} {expr.op} {right_s}"

    if isinstance(expr, FunctionCall):
        args = ", ".join(format_expr(a) for a in expr.args)
        return f"{expr.name}({args})"

    if isinstance(expr, FormattedString):
        return "f" + _escape_string(expr.template)

    if isinstance(expr, Assignment):
        return f"{expr.name} = {format_expr(expr.value)}"

    if isinstance(expr, IndexAssignment):
        return f"{format_expr(expr.target)}[{format_expr(expr.index)}] = {format_expr(expr.value)}"

    raise TypeError(f"Unsupported expression node: {type(expr).__name__}")


def _indent(lines: list[str], level: int) -> list[str]:
    prefix = " " * (4 * level)
    return [prefix + line if line else line for line in lines]


def format_stmt(stmt, level: int) -> list[str]:
    if isinstance(stmt, Assignment):
        return _indent([format_expr(stmt)], level)

    if isinstance(stmt, IndexAssignment):
        return _indent([format_expr(stmt)], level)

    if isinstance(stmt, FunctionCall):
        return _indent([format_expr(stmt)], level)

    if isinstance(stmt, ReturnStatement):
        return _indent([f"parat {format_expr(stmt.value)}"], level)

    if isinstance(stmt, BreakStatement):
        return _indent(["todu"], level)

    if isinstance(stmt, ContinueStatement):
        return _indent(["pudhe"], level)

    if isinstance(stmt, ThrowStatement):
        return _indent([f"feka {format_expr(stmt.value)}"], level)

    if isinstance(stmt, FunctionDeclaration):
        header = f"kam {stmt.name}({', '.join(stmt.params)}) " + "{"
        lines = _indent([header], level)
        for inner in stmt.body:
            lines.extend(format_stmt(inner, level + 1))
        lines.extend(_indent(["}"], level))
        return lines

    if isinstance(stmt, IfStatement):
        lines = _indent([f"jar {format_expr(stmt.condition)} " + "{"], level)
        for inner in stmt.body:
            lines.extend(format_stmt(inner, level + 1))
        if stmt.else_body is not None:
            lines.extend(_indent(["}"], level))
            lines.extend(_indent(["nahitar {"], level))
            for inner in stmt.else_body:
                lines.extend(format_stmt(inner, level + 1))
            lines.extend(_indent(["}"], level))
        else:
            lines.extend(_indent(["}"], level))
        return lines

    if isinstance(stmt, WhileStatement):
        lines = _indent([f"joparyant {format_expr(stmt.condition)} " + "{"], level)
        for inner in stmt.body:
            lines.extend(format_stmt(inner, level + 1))
        lines.extend(_indent(["}"], level))
        return lines

    if isinstance(stmt, ForLoop):
        lines = _indent([f"sathi {format_expr(stmt.condition)} " + "{"], level)
        for inner in stmt.body:
            lines.extend(format_stmt(inner, level + 1))
        lines.extend(_indent(["}"], level))
        return lines

    if isinstance(stmt, TryStatement):
        lines = _indent(["prayatna {"], level)
        for inner in stmt.try_body:
            lines.extend(format_stmt(inner, level + 1))

        if stmt.catch_body is not None:
            lines.extend(_indent(["}"], level))
            if stmt.catch_name:
                lines.extend(_indent([f"pakad ({stmt.catch_name}) " + "{"], level))
            else:
                lines.extend(_indent(["pakad {"], level))
            for inner in stmt.catch_body:
                lines.extend(format_stmt(inner, level + 1))

        if stmt.finally_body is not None:
            lines.extend(_indent(["}"], level))
            lines.extend(_indent(["shevati {"], level))
            for inner in stmt.finally_body:
                lines.extend(format_stmt(inner, level + 1))

        lines.extend(_indent(["}"], level))
        return lines

    if stmt.__class__.__name__ == "PrintStatement":
        # avoid import cycle: PrintStatement is defined in parser.py
        args = ", ".join(format_expr(a) for a in stmt.args)
        return _indent([f"dakhav({args})"], level)

    raise TypeError(f"Unsupported statement node: {type(stmt).__name__}")


def format_program(program: Program) -> str:
    lines: list[str] = []
    for stmt in program.statements:
        lines.extend(format_stmt(stmt, 0))
    return "\n".join(lines).rstrip() + "\n"


def format_source(source: str) -> str:
    tokens = Lexer(source).tokenize()
    program = Parser(tokens, source=source).parse()
    return format_program(program)
