"""
Nexa Language — Flamegraph Profiler v24
nexa profile file.nexa --output profile.json
Generates speedscope-compatible JSON flamegraph.
"""
from __future__ import annotations
import time
import json
from typing import List, Dict, Any


class CallFrame:
    __slots__ = ("name", "t_start", "children", "duration")
    def __init__(self, name: str):
        self.name = name
        self.t_start = time.perf_counter()
        self.children: List["CallFrame"] = []
        self.duration = 0.0

    def finish(self):
        self.duration = (time.perf_counter() - self.t_start) * 1000


class NexaProfiler:
    def __init__(self):
        self._stack: List[CallFrame] = []
        self._root = CallFrame("<root>")
        self._stack.append(self._root)

    def enter(self, name: str):
        frame = CallFrame(name)
        self._stack[-1].children.append(frame)
        self._stack.append(frame)

    def leave(self):
        if len(self._stack) > 1:
            frame = self._stack.pop()
            frame.finish()

    def _frame_to_dict(self, frame: CallFrame) -> dict:
        return {
            "name": frame.name,
            "duration_ms": round(frame.duration, 4),
            "children": [self._frame_to_dict(c) for c in frame.children]
        }

    def to_speedscope(self) -> dict:
        """Export as speedscope.app-compatible JSON."""
        frames = []
        samples = []
        weights = []

        def _walk(frame: CallFrame, stack: list):
            idx = len(frames)
            frames.append({"name": frame.name})
            samples.append(stack + [idx])
            weights.append(frame.duration)
            for child in frame.children:
                _walk(child, stack + [idx])

        for child in self._root.children:
            _walk(child, [])

        return {
            "$schema": "https://www.speedscope.app/file-format-schema.json",
            "shared": {"frames": frames},
            "profiles": [{
                "type": "sampled",
                "name": "Nexa Profile",
                "unit": "milliseconds",
                "startValue": 0,
                "endValue": sum(weights),
                "samples": samples,
                "weights": weights
            }]
        }

    def report(self):
        """Print a simple text report."""
        print("\n  🔥 Nexa Profile Report")
        print(f"  {'─' * 50}")

        def _print(frame: CallFrame, indent: int = 0):
            bar = "  " * indent + ("├─ " if indent else "")
            print(f"  {bar}{frame.name:<35} {frame.duration:8.3f} ms")
            for c in sorted(frame.children, key=lambda x: -x.duration):
                _print(c, indent + 1)

        for child in sorted(self._root.children, key=lambda x: -x.duration):
            _print(child)
        print(f"  {'─' * 50}")
        total = sum(c.duration for c in self._root.children)
        print(f"  Total: {total:.3f} ms")


def run_profile(source_path: str, output_json: str = None):
    """Profile a Nexa source file and output results."""
    from .lexer import Lexer
    from .parser import Parser
    from .interpreter import Interpreter

    src = open(source_path, encoding="utf-8").read()
    tokens = Lexer(src).tokenize()
    ast = Parser(tokens).parse()

    profiler = NexaProfiler()

    # Monkey-patch interpreter to record exec calls
    interp = Interpreter()
    original_exec = interp._exec.__func__ if hasattr(interp._exec, '__func__') else None

    def profiled_exec(node, env):
        name = type(node).__name__
        profiler.enter(name)
        try:
            return interp._exec.__func__(interp, node, env)
        finally:
            profiler.leave()

    interp._exec = profiled_exec

    t0 = time.perf_counter()
    interp.run(ast)
    total = (time.perf_counter() - t0) * 1000

    profiler.report()

    if output_json:
        data = profiler.to_speedscope()
        with open(output_json, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2)
        print(f"\n  📄 Flamegraph JSON saved → {output_json}")
        print(f"  🌐 Open at: https://www.speedscope.app (drag & drop the file)")
    else:
        print(f"\n  💡 Tip: use --output profile.json to export a speedscope flamegraph")
