"""
Nexa Language — Python Transpiler
`nexa compile file.nexa` → generates `file.py`

Walks the Nexa AST and emits idiomatic Python 3.9+ code.
"""
from __future__ import annotations
import os
import re
from typing import Any, List
from .ast_nodes import *


class Transpiler:
    """Walks a Nexa AST and emits Python source code."""

    INDENT = "    "

    def __init__(self):
        self._depth = 0
        self._var_counter = 0
        self._tests = []
        self._scopes = [set()]
        self._nonlocals = [set()]
        self._globals = [set()]

    # ── Public API ────────────────────────────────────────────────────────
    def transpile(self, program: Program) -> str:
        header = (
            "# ── Generated by Nexa Transpiler ──────────────────────────────\n"
            "# This file was auto-generated from a .nexa source file.\n"
            "# Do not edit directly — regenerate with: nexa compile <file>\n"
            "# ─────────────────────────────────────────────────────────────\n"
            "from __future__ import annotations\n"
            "import asyncio, math, json, os, random, functools, time as _time\n"
            "import sys as _sys\n"
            "_sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))\n"
            "_sys.path.insert(0, os.path.expanduser('~/.nexa/packages'))\n"
            "from nexa_lang.builtins import BUILTINS\n"
            "_g = globals()\n"
            "for _k, _v in BUILTINS.items():\n"
            "    if _k not in ('true', 'false', 'null'): _g[_k] = _v\n"
            "true = True; false = False; null = None\n\n"
            "# ── Nexa runtime stubs ───────────────────────────────────────\n"
            "_nexa_modules = {}\n"
            "def _nexa_opt_access(obj, attr):\n"
            "    if obj is None: return None\n"
            "    if isinstance(obj, dict): return obj.get(attr)\n"
            "    return getattr(obj, attr, None)\n\n"
            "def _nexa_invoke_method(obj, method, *args, **kwargs):\n"
            "    if isinstance(obj, list):\n"
            "        if method == 'push': return obj.append(*args, **kwargs)\n"
            "        if method == 'join': return (args[0] if args else '').join(str(x) for x in obj)\n"
            "        if method == 'map': return list(map(args[0], obj))\n"
            "        if method == 'filter': return list(filter(args[0], obj))\n"
            "        if method == 'reduce': return functools.reduce(args[0], obj, *args[1:]) if len(args) > 1 else functools.reduce(args[0], obj)\n"
            "    if isinstance(obj, dict) and method == 'keys': return list(obj.keys())\n"
            "    if isinstance(obj, dict) and method == 'values': return list(obj.values())\n"
            "    if str(type(obj)) == '<class \\'super\\'>' and method == 'init': method = '__init__'\n"
            "    return getattr(obj, method)(*args, **kwargs)\n\n"
        )
        body = "\n".join(self._stmt(s) for s in program.statements)
        if self._tests:
            body += "\nif __name__ == '__main__':\n"
            body += "    print('\\n========================================')\n"
            body += "    print('Nexa Transpiled Test Runner (`nexa compile`)')\n"
            body += "    print('========================================\\n')\n"
            for t in self._tests:
                body += f"    print('  🧪 Running: \"{t}\"')\n"
                body += f"    test_{t}()\n"
                body += f"    print('  \\033[92m✅ PASS: {t}\\033[0m\\n')\n"
            body += "    print('All tests passed natively via transpilation!\\n')\n"
        return header + body + "\n"

    # ── Indentation helpers ────────────────────────────────────────────────
    def _ind(self) -> str:   return self.INDENT * self._depth
    def _fresh(self, prefix: str = "var") -> str:
        self._var_counter += 1
        return f"__{prefix}_{self._var_counter}"
    def _in(self):           self._depth += 1
    def _out(self):          self._depth -= 1

    def _block(self, stmts: list) -> str:
        if not stmts: return self._ind() + "    pass"
        self._in()
        lines = "\n".join(self._stmt(s) for s in stmts)
        self._out()
        return lines

    # ── Statements ────────────────────────────────────────────────────────
    def _stmt(self, node: Any) -> str:
        method = f"_s_{type(node).__name__}"
        handler = getattr(self, method, None)
        if handler: return handler(node)
        return f"{self._ind()}{self._expr(node)}"

    def _s_FunctionDef(self, node: FunctionDef) -> str:
        self._scopes[-1].add(node.name)
        params = []
        clean_params = []
        is_static = any(getattr(d.expr, "name", "") == "staticmethod" for d in getattr(node, "decorators", []))
        if getattr(self, "in_class", False) and not is_static and (not node.params or node.params[0] != "self"):
            params.append("self")
            clean_params.append("self")
        for p in node.params:
            if p.startswith("..."):
                clean_params.append(p[3:])
                params.append(f"*{p[3:]}")
                continue
            clean_params.append(p)
            ann = f": {node.param_types[p]}" if hasattr(node, 'param_types') and p in node.param_types else ""
            default = f" = {self._expr(node.defaults[p])}" if hasattr(node, 'defaults') and p in node.defaults else ""
            params.append(f"{p}{ann}{default}")
        ret = f" -> {node.return_type}" if getattr(node, 'return_type', None) else ""
        prefix = "async " if getattr(node, 'is_async', False) else ""
        
        decorators = "\n".join(f"{self._ind()}@{self._expr(d.expr)}" for d in node.decorators)
        if decorators: decorators += "\n"
        
        name = "__init__" if getattr(self, "in_class", False) and node.name == "init" else node.name
        sig = f"{decorators}{self._ind()}{prefix}def {name}({', '.join(params)}){ret}:"
        doc = ""
        if getattr(node, 'docstring', None):
            self._in()
            doc = f'\n{self._ind()}"""{node.docstring}"""'
            self._out()
            
        self._scopes.append(set(clean_params))
        self._nonlocals.append(set())
        self._globals.append(set())
            
        body = self._block(node.body)
        
        nonlocals = self._nonlocals.pop()
        globals_ = self._globals.pop()
        self._scopes.pop()
        
        injection = ""
        if globals_:
            for g in globals_: injection += f"{self._ind()}    global {g}\n"
        if nonlocals:
            for nl in nonlocals: injection += f"{self._ind()}    nonlocal {nl}\n"
            
        return f"{sig}{doc}\n{injection}{body}"

    def _s_ClassDef(self, node: ClassDef) -> str:
        base = f"({node.base})" if node.base else ""
        
        decorators = "\n".join(f"{self._ind()}@{self._expr(d.expr)}" for d in node.decorators)
        if decorators: decorators += "\n"
        
        header = f"{decorators}{self._ind()}class {node.name}{base}:"
        
        prev_class = getattr(self, "in_class", False)
        self.in_class = True
        body = self._block(node.body)
        self.in_class = prev_class
        
        return f"{header}\n{body}"

    def _s_StructDef(self, node: StructDef) -> str:
        fields = ", ".join(
            f"{name}: {hint or 'object'}" for name, hint in node.fields
        )
        defaults = ", ".join(f"self.{name} = {name}" for name, _ in node.fields)
        params = ", ".join(name for name, _ in node.fields)
        return (
            f"{self._ind()}class {node.name}:\n"
            f"{self._ind()}    def __init__(self, {params}):\n"
            + "".join(f"{self._ind()}        self.{n} = {n}\n" for n, _ in node.fields)
            + f"{self._ind()}    def __repr__(self):\n"
            f"{self._ind()}        return f\"{node.name}({', '.join(n+'={self.'+n+'}' for n, _ in node.fields)})\"\n"
        )

    def _s_EnumDef(self, node: EnumDef) -> str:
        lines = [
            f"{self._ind()}from enum import Enum as _Enum",
            f"{self._ind()}class {node.name}(_Enum):"
        ]
        self._in()
        for i, m in enumerate(node.members, 1):
            lines.append(f"{self._ind()}{m} = {i}")
        self._out()
        return "\n".join(lines)

    def _s_LetStatement(self, node: LetStatement) -> str:
        self._scopes[-1].add(node.name)
        ann = f": {node.type_hint}" if getattr(node, 'type_hint', None) else ""
        return f"{self._ind()}{node.name}{ann} = {self._expr(node.value)}"

    def _s_DestructuredAssignment(self, node: DestructuredAssignment) -> str:
        if node.is_dict:
            tmp = f"_tmp_{id(node)}"
            lines = [f"{self._ind()}{tmp} = {self._expr(node.value)}"]
            for target in node.targets:
                lines.append(f"{self._ind()}{target} = {tmp}.get('{target}')")
            return "\n".join(lines)
        else:
            targets = ", ".join(node.targets)
            return f"{self._ind()}{targets} = {self._expr(node.value)}"

    def _s_ReturnStatement(self, node: ReturnStatement) -> str:
        return f"{self._ind()}return {self._expr(node.value)}"

    def _s_YieldStatement(self, node: YieldStatement) -> str:
        return f"{self._ind()}yield {self._expr(node.value)}"

    def _s_ExprStatement(self, node: ExprStatement) -> str:
        return f"{self._ind()}{self._expr(node.expr)}"

    def _e_IfExpr(self, node: IfExpr) -> str:
        cond = self._expr(node.condition)
        # For transpiling implicit returns cleanly, we should support statement vs expr transpilation.
        # But for Python, IfExpr as a statement:
        result = f"{self._ind()}if {cond}:\n{self._block(node.then_block)}"
        if node.else_block:
            # Check if else_block is a single IfExpr (else if chain)
            if len(node.else_block) == 1 and isinstance(node.else_block[0], IfExpr):
                inner = self._e_IfExpr(node.else_block[0])
                inner = self._ind() + "el" + inner.lstrip()
                result += f"\n{inner}"
            else:
                result += f"\n{self._ind()}else:\n{self._block(node.else_block)}"
        return result

    def _s_ForStatement(self, node: ForStatement) -> str:
        return (f"{self._ind()}for {node.var} in {self._expr(node.iterable)}:\n"
                f"{self._block(node.body)}")

    def _s_WhileStatement(self, node: WhileStatement) -> str:
        return (f"{self._ind()}while {self._expr(node.condition)}:\n"
                f"{self._block(node.body)}")

    def _s_BreakStatement(self, node: BreakStatement) -> str:
        return f"{self._ind()}break"

    def _s_ContinueStatement(self, node: ContinueStatement) -> str:
        return f"{self._ind()}continue"

    def _s_InterfaceDef(self, node) -> str:
        # Interfaces are a compile-time concept; emit a comment
        methods = ', '.join(node.method_names)
        return f"{self._ind()}# interface {node.name} {{{methods}}}"

    def _s_RangeExpr(self, node) -> str:
        # Used as a statement (rare)
        return f"{self._ind()}{self._e_RangeExpr(node)}"

    def _e_RangeExpr(self, node) -> str:
        s = self._expr(node.start)
        e = self._expr(node.end)
        return f"list(range({s}, {e} + 1))" if node.inclusive else f"list(range({s}, {e}))"

    def _e_TryExpr(self, node: TryExpr) -> str:
        result = f"{self._ind()}try:\n{self._block(node.try_block)}"
        if node.catch_block:
            bind = f" as {node.catch_bind}" if getattr(node, 'catch_bind', None) else ""
            result += (f"\n{self._ind()}except Exception{bind}:\n"
                       f"{self._block(node.catch_block)}")
        if node.finally_block:
            result += (f"\n{self._ind()}finally:\n"
                       f"{self._block(node.finally_block)}")
        return result

    def _s_ThrowStatement(self, node: ThrowStatement) -> str:
        return f"{self._ind()}raise Exception({self._expr(node.value)})"

    def _e_MatchExpr(self, node: MatchExpr) -> str:
        # Evaluate subject once to aviod re-triggering side effects
        subject = self._expr(node.subject)
        var_name = self._fresh("match")
        lines = [f"{self._ind()}{var_name} = {subject}"]
        
        first = True
        for case in node.cases:
            cond = f"{var_name} == {self._expr(case.pattern)}"
            kw = "if" if first else "elif"
            lines.append(f"{self._ind()}{kw} {cond}:")
            self._in()
            lines.append(self._block(case.body))
            self._out()
            first = False
            
        if node.default_body:
            if first:
                lines.append(self._block(node.default_body))
            else:
                lines.append(f"{self._ind()}else:")
                self._in()
                lines.append(self._block(node.default_body))
                self._out()
        return "\n".join(lines)

    def _s_ImportStatement(self, node: ImportStatement) -> str:
        if node.is_file:
            alias = node.alias or os.path.splitext(os.path.basename(node.module))[0]
            return f"{self._ind()}# import \"{node.module}\" as {alias}  (Nexa module)"
            
        known = {"automl","nlp","vision","math","stats","fs","json",
                 "http","time","random","agent","df",
                 "db", "re", "serve", "tensor", "env", "sys",
                 "datetime", "crypto", "base64", "task"}
        if node.module in known:
            alias = node.alias or node.module
            return f"{self._ind()}{alias} = BUILTINS['{node.module}']"
            
        return f"{self._ind()}import {node.module}" + (f" as {node.alias}" if node.alias else "")

    def _s_FromImportStatement(self, node: FromImportStatement) -> str:
        names = ", ".join(node.names)
        mod = node.module[:-5] if node.module.endswith(".nexa") else node.module
        return f"{self._ind()}from {mod} import {names}"

    def _s_ModelDef(self, node: ModelDef) -> str:
        return f"{self._ind()}{node.name} = None  # NexaModel '{node.name}'"

    def _s_TestBlock(self, node: TestBlock) -> str:
        name = re.sub(r'[^a-zA-Z0-9_]', '_', node.name)
        self._tests.append(name)
        fn_name = "test_" + name
        self._scopes[-1].add(fn_name)
        
        self._scopes.append(set())
        self._nonlocals.append(set())
        self._globals.append(set())
        
        header = f"{self._ind()}def {fn_name}():"
        body = self._block(node.body)
        
        nonlocals = self._nonlocals.pop()
        globals_ = self._globals.pop()
        self._scopes.pop()
        
        injection = ""
        if globals_:
            for g in globals_: injection += f"{self._ind()}    global {g}\n"
        if nonlocals:
            for nl in nonlocals: injection += f"{self._ind()}    nonlocal {nl}\n"
            
        return f"{header}\n{injection}{body}"

    def _s_AssertStatement(self, node: AssertStatement) -> str:
        args = [self._expr(a) for a in node.args]
        kind = node.kind
        if kind == "eq":      return f"{self._ind()}assert {args[0]} == {args[1]}"
        if kind == "ne":      return f"{self._ind()}assert {args[0]} != {args[1]}"
        if kind == "true":    return f"{self._ind()}assert {args[0]}"
        if kind == "false":   return f"{self._ind()}assert not {args[0]}"
        if kind == "contains":return f"{self._ind()}assert {args[1]} in {args[0]}"
        return f"{self._ind()}pass  # assert_raises not yet supported"

    # ── Expressions ───────────────────────────────────────────────────────
    def _expr(self, node: Any) -> str:
        if isinstance(node, NumberLiteral):  return repr(node.value)
        if isinstance(node, StringLiteral):  return repr(node.value)
        if isinstance(node, BoolLiteral):    return "True" if node.value else "False"
        if isinstance(node, NullLiteral):    return "None"
        if isinstance(node, SelfRef):        return "self"
        if isinstance(node, SuperRef):       return "super()"
        if isinstance(node, FString):        return self._expr_fstring(node)
        if isinstance(node, Identifier):     return node.name
        if isinstance(node, ListLiteral):
            return "[" + ", ".join(f"*{self._expr(e.expr)}" if isinstance(e, SpreadExpr) else self._expr(e) for e in node.elements) + "]"
        if isinstance(node, SetLiteral):
            if not node.elements: return "set()"
            return "{" + ", ".join(f"*{self._expr(e.expr)}" if isinstance(e, SpreadExpr) else self._expr(e) for e in node.elements) + "}"
        if isinstance(node, DictLiteral):
            pairs = []
            for k, v in node.pairs:
                if isinstance(k, SpreadExpr): pairs.append(f"**{self._expr(k.expr)}")
                else: pairs.append(f"{self._expr(k)}: {self._expr(v)}")
            return "{" + ", ".join(pairs) + "}"
        if isinstance(node, ListComp):       return self._expr_listcomp(node)
        if isinstance(node, RangeExpr):      return self._e_RangeExpr(node)
        if isinstance(node, BinaryOp):
            op = "and" if node.op == "and" else "or" if node.op == "or" else node.op
            return f"({self._expr(node.left)} {op} {self._expr(node.right)})"
        if isinstance(node, UnaryOp):
            op = node.op if node.op != "not" else "not "
            return f"({op}{self._expr(node.operand)})"
        if isinstance(node, TernaryOp):
            return f"({self._expr(node.true_expr)} if {self._expr(node.condition)} else {self._expr(node.false_expr)})"
        if isinstance(node, NullishCoalescing):
            return f"({self._expr(node.left)} if {self._expr(node.left)} is not None else {self._expr(node.right)})"
        if isinstance(node, OptionalAccess):
            return f"_nexa_opt_access({self._expr(node.obj)}, '{node.member}')"
        if isinstance(node, PipelineExpr):   return self._expr_pipeline(node)
        if isinstance(node, Assignment):
            name = node.name
            if len(self._scopes) > 1 and name not in self._scopes[-1]:
                is_global = False; found = False
                for i in range(len(self._scopes)-2, -1, -1):
                    if name in self._scopes[i]:
                        found = True; is_global = (i == 0); break
                if found:
                    if is_global: self._globals[-1].add(name)
                    else: self._nonlocals[-1].add(name)
            return f"{node.name} = {self._expr(node.value)}"
        if isinstance(node, CompoundAssignment):
            name = node.name
            if len(self._scopes) > 1 and name not in self._scopes[-1]:
                is_global = False; found = False
                for i in range(len(self._scopes)-2, -1, -1):
                    if name in self._scopes[i]:
                        found = True; is_global = (i == 0); break
                if found:
                    if is_global: self._globals[-1].add(name)
                    else: self._nonlocals[-1].add(name)
            return f"{node.name} {node.op}= {self._expr(node.value)}"
        if isinstance(node, MemberAssignment):
            return f"{self._expr(node.obj)}.{node.member} = {self._expr(node.value)}"
        if isinstance(node, IndexAssignment):
            return f"{self._expr(node.obj)}[{self._expr(node.index)}] = {self._expr(node.value)}"
        if isinstance(node, MemberAccess):
            return f"{self._expr(node.obj)}.{node.member}"
        if isinstance(node, IndexAccess):
            if type(node.index).__name__ == "SliceExpr":
                start = self._expr(node.index.start) if node.index.start else ""
                stop = self._expr(node.index.stop) if node.index.stop else ""
                step = f":{self._expr(node.index.step)}" if node.index.step else ""
                return f"{self._expr(node.obj)}[{start}:{stop}{step}]"
            return f"{self._expr(node.obj)}[{self._expr(node.index)}]"
        if isinstance(node, FunctionCall):   return self._expr_call(node)
        if isinstance(node, AwaitExpr):
            expr = self._expr(node.expr)
            return f"asyncio.run({expr})" if self._depth == 0 else f"await {expr}"
        if isinstance(node, NewExpr):
            args = ", ".join(self._expr(a) for a in node.args)
            kwargs = ", ".join(f"{k}={self._expr(v)}" for k, v in node.kwargs.items())
            all_args = ", ".join(filter(None, [args, kwargs]))
            return f"{node.class_name}({all_args})"
        if isinstance(node, IfExpr):         return self._e_IfExpr(node)
        if isinstance(node, TryExpr):        return self._e_TryExpr(node)
        if isinstance(node, MatchExpr):      return self._e_MatchExpr(node)
        if isinstance(node, BlockExpr):      return f"(\n{self._block(node.stmts)}\n{self._ind()})"
        
        if isinstance(node, LambdaExpr):
            if len(node.body) == 1 and isinstance(node.body[0], ReturnStatement):
                return f"lambda {', '.join(node.params)}: {self._expr(node.body[0].value)}"
            else:
                raise Exception("Multi-line lambdas not supported in transpiler yet")
        return repr(node)

    def _expr_fstring(self, node: FString) -> str:
        parts = []
        for kind, val in node.parts:
            if kind == "lit": parts.append(val.replace("{", "{{").replace("}", "}}"))
            else: parts.append("{" + self._expr(val) + "}")
        return 'f"' + "".join(parts) + '"'

    def _expr_listcomp(self, node: ListComp) -> str:
        cond = f" if {self._expr(node.condition)}" if node.condition else ""
        return (f"[{self._expr(node.expr)} for {node.var} in "
                f"{self._expr(node.iterable)}{cond}]")

    def _expr_pipeline(self, node: PipelineExpr) -> str:
        # a |> f |> g  →  g(f(a))
        value = self._expr(node.stages[0])
        for stage in node.stages[1:]:
            fn = self._expr(stage)
            value = f"{fn}({value})"
        return value

    def _expr_call(self, node: FunctionCall) -> str:
        # Check for list.map(), list.filter(), list.reduce() interceptions
        if isinstance(node.callee, MemberAccess) and node.callee.member in ("map", "filter", "reduce"):
            obj_str = self._expr(node.callee.obj)
            fn_str = self._expr(node.args[0])
            if node.callee.member == "map":
                return f"list(map({fn_str}, {obj_str}))"
            if node.callee.member == "filter":
                return f"list(filter({fn_str}, {obj_str}))"
            if node.callee.member == "reduce":
                args_rest = ", ".join(self._expr(a) for a in node.args[1:])
                if args_rest:
                    return f"functools.reduce({fn_str}, {obj_str}, {args_rest})"
                return f"functools.reduce({fn_str}, {obj_str})"

        args_list = []
        for a in node.args:
            if isinstance(a, SpreadExpr): args_list.append(f"*{self._expr(a.expr)}")
            else: args_list.append(self._expr(a))
        args = ", ".join(args_list)
        kwargs = ", ".join(f"{k}={self._expr(v)}" for k, v in node.kwargs.items())
        all_args = ", ".join(filter(None, [args, kwargs]))
        
        # Intercept native method calls obj.method() routing them through _nexa_invoke_method
        if isinstance(node.callee, MemberAccess):
            obj = self._expr(node.callee.obj)
            method = node.callee.member
            return f"_nexa_invoke_method({obj}, '{method}'" + (f", {all_args})" if all_args else ")")
            
        callee = self._expr(node.callee)
        return f"{callee}({all_args})"

import os  # needed for ImportStatement transpilation
