#!/usr/bin/env python3
"""
pybreakz - A CLI debugger for Claude Code
Run Python scripts with breakpoints and get structured output.

Usage:
    pybreakz run script.py --breakpoints 10,25,42
    pybreakz run script.py --breakpoints 42 --watch "x,result"
    pybreakz run script.py --on-exception
    pybreakz run script.py --breakpoints "42:x>10" --eval "len(items),type(r)"
    pybreakz run script.py --breakpoints 10 -- --my-arg foo --other-arg 42

    # Cross-file breakpoints (set breakpoints in imported modules):
    pybreakz run script.py --breakpoints "utils.py:55,helpers.py:12:x>0"

    # Variable watchpoints (break whenever a variable changes):
    pybreakz run script.py --trace-vars "counter,result"
    pybreakz run script.py --breakpoints 10 --trace-vars "x"
"""

import sys
import os
import ast
import json
import argparse
import traceback
import threading
import signal
import textwrap
from types import FrameType
from typing import Any, Optional

_UNSET = object()  # sentinel: variable not yet seen


# ─── Output formatting ───────────────────────────────────────────────────────

SEPARATOR = "─" * 60

def fmt_value(v: Any, max_len: int = 200) -> str:
    """Format a value for display, truncating if too long."""
    try:
        if isinstance(v, str):
            r = repr(v)
        elif isinstance(v, (list, tuple, set)):
            r = f"{type(v).__name__}[{len(v)} items]: {repr(v)}"
        elif isinstance(v, dict):
            r = f"dict[{len(v)} keys]: {repr(v)}"
        else:
            r = repr(v)
        if len(r) > max_len:
            return r[:max_len] + "  ..."
        return r
    except Exception as e:
        return f"<error formatting value: {e}>"


def fmt_frame(frame: FrameType) -> list[str]:
    """Format a call stack from a frame, filtering out pybreakz internals."""
    _self = os.path.abspath(__file__)
    stack = []
    f = frame
    while f is not None:
        ffile = os.path.abspath(f.f_code.co_filename)
        if ffile != _self:  # skip pybreakz's own frames
            fname = os.path.basename(ffile)
            stack.append(f"  {f.f_code.co_name}()  [{fname}:{f.f_lineno}]")
        f = f.f_back
    stack.reverse()
    return stack


def fmt_locals(local_vars: dict, watch: list[str], eval_exprs: list[str], frame: FrameType) -> list[str]:
    """Format locals, watched vars, and eval expressions."""
    lines = []

    # Filter out dunder/internal vars
    visible = {k: v for k, v in local_vars.items() if not k.startswith("__")}

    if visible:
        lines.append("Locals:")
        for k, v in visible.items():
            lines.append(f"  {k:<20} = {fmt_value(v)}")
    else:
        lines.append("Locals: (none)")

    if watch:
        lines.append("Watched:")
        combined = {**frame.f_globals, **local_vars}
        for expr in watch:
            try:
                result = eval(expr, combined)
                lines.append(f"  {expr:<20} = {fmt_value(result)}")
            except Exception as e:
                lines.append(f"  {expr:<20} ! {e}")

    if eval_exprs:
        lines.append("Evaluated:")
        combined = {**frame.f_globals, **local_vars}
        for expr in eval_exprs:
            try:
                result = eval(expr, combined)
                lines.append(f"  {expr:<20} = {fmt_value(result)}")
            except Exception as e:
                lines.append(f"  {expr:<20} ! {e}")

    return lines


# ─── Breakpoint parsing ───────────────────────────────────────────────────────

def parse_breakpoints(bp_str: str, script_basename: str) -> dict[str, dict[int, Optional[str]]]:
    """
    Parse breakpoint specs. Supports:
      "42"                 → line 42 in the main script
      "42:x>10"            → conditional, main script
      "utils.py:42"        → line 42 in utils.py (any imported file)
      "utils.py:42:x>10"   → conditional, in utils.py

    Returns dict of {basename: {line: condition_or_None}}
    """
    result: dict[str, dict[int, Optional[str]]] = {}
    for part in bp_str.split(","):
        part = part.strip()
        if not part:
            continue
        first, _, rest = part.partition(":")
        try:
            # First segment is a plain integer → main-script line number
            line = int(first)
            filename = script_basename
            condition = rest.strip() if rest else None
        except ValueError:
            # First segment is a filename
            filename = first.strip()
            if not rest:
                raise ValueError(
                    f"Expected 'file:line' or 'file:line:condition', got: {part!r}"
                )
            line_s, _, condition_s = rest.partition(":")
            try:
                line = int(line_s.strip())
            except ValueError:
                raise ValueError(f"Invalid line number in breakpoint: {part!r}")
            condition = condition_s.strip() if condition_s else None

        result.setdefault(filename, {})[line] = condition
    return result


# ─── Tracer ──────────────────────────────────────────────────────────────────

class Debugger:
    def __init__(
        self,
        script_path: str,
        breakpoints: dict[str, dict[int, Optional[str]]],
        watch: list[str],
        eval_exprs: list[str],
        on_exception: bool,
        timeout: int,
        output_format: str,
        max_hits: int,
        trace_vars: list[str],
    ):
        self.script_path = os.path.abspath(script_path)
        self.breakpoints = breakpoints      # {basename: {line: condition}}
        self.watch = watch
        self.eval_exprs = eval_exprs
        self.on_exception = on_exception
        self.timeout = timeout
        self.output_format = output_format
        self.max_hits = max_hits
        self.trace_vars = trace_vars

        self.hits: list[dict] = []          # breakpoint hits + watchpoint hits, in order
        self.hit_count = 0
        self.exception_info: Optional[dict] = None
        self._var_last: dict[tuple[int, str], Any] = {}   # (code_id, var) → last value
        self._prev_line: dict[int, int] = {}              # code_id → previous line number

    def _check_condition(self, condition: str, frame: FrameType) -> bool:
        try:
            combined = {**frame.f_globals, **frame.f_locals}
            return bool(eval(condition, combined))
        except Exception:
            return False  # Skip if condition errors

    def _check_var_changes(self, frame: FrameType, basename: str, lineno: int, filepath: str):
        """Record a watchpoint hit for any traced variable that changed value."""
        code_id = id(frame.f_code)
        combined = {**frame.f_globals, **frame.f_locals}
        for var in self.trace_vars:
            key = (code_id, var)
            current = combined.get(var, _UNSET)
            last = self._var_last.get(key, _UNSET)

            if last is not _UNSET and current is not _UNSET:
                # Compare; handle objects like numpy arrays that return non-bool
                try:
                    changed = current != last
                    if not isinstance(changed, bool):
                        changed = bool(changed)
                except Exception:
                    changed = id(current) != id(last)

                if changed and self.hit_count < self.max_hits:
                    self._record_watchpoint(frame, var, last, current, basename, lineno, filepath)

            if current is not _UNSET:
                self._var_last[key] = current

    def trace_calls(self, frame: FrameType, event: str, arg: Any):
        """sys.settrace handler — processes line events in all files."""
        if event != "line":
            return self.trace_calls

        filepath = os.path.abspath(frame.f_code.co_filename)
        basename = os.path.basename(filepath)
        lineno = frame.f_lineno

        code_id = id(frame.f_code)

        # 1. Variable watchpoints: attribute changes to the line that just ran
        #    (the tracer fires at the START of each line, so changes from the
        #    previous line are visible now — report the previous line number)
        if self.trace_vars:
            prev_lineno = self._prev_line.get(code_id, lineno)
            self._check_var_changes(frame, basename, prev_lineno, filepath)
        self._prev_line[code_id] = lineno

        # 2. Line breakpoints (only files with registered breakpoints)
        file_bps = self.breakpoints.get(basename, {})
        if lineno in file_bps and self.hit_count < self.max_hits:
            condition = file_bps[lineno]
            if condition is None or self._check_condition(condition, frame):
                self._record_hit(frame, lineno, filepath, basename, condition)
                self.hit_count += 1

        return self.trace_calls

    def _record_hit(self, frame: FrameType, lineno: int, filepath: str, basename: str, condition: Optional[str]):
        stack = fmt_frame(frame)
        local_lines = fmt_locals(
            dict(frame.f_locals), self.watch, self.eval_exprs, frame
        )
        self.hits.append({
            "type": "breakpoint",
            "file": basename,
            "line": lineno,
            "condition": condition,
            "stack": stack,
            "locals": local_lines,
            "source_line": self._get_source_line(lineno, filepath),
        })

    def _record_watchpoint(
        self,
        frame: FrameType,
        var: str,
        old_val: Any,
        new_val: Any,
        basename: str,
        lineno: int,
        filepath: str,
    ):
        stack = fmt_frame(frame)
        self.hits.append({
            "type": "watchpoint",
            "var": var,
            "old": fmt_value(old_val),
            "new": fmt_value(new_val),
            "file": basename,
            "line": lineno,
            "source_line": self._get_source_line(lineno, filepath),
            "stack": stack,
        })
        self.hit_count += 1

    def _get_source_line(self, lineno: int, filepath: str = None) -> str:
        path = filepath or self.script_path
        try:
            with open(path) as f:
                lines = f.readlines()
            if 1 <= lineno <= len(lines):
                return lines[lineno - 1].rstrip()
        except Exception:
            pass
        return ""

    def _get_context_lines(self, lineno: int, context: int = 3) -> list[tuple[int, str, bool]]:
        """Get source lines around a line number. Returns (lineno, text, is_target)."""
        try:
            with open(self.script_path) as f:
                lines = f.readlines()
            result = []
            start = max(1, lineno - context)
            end = min(len(lines), lineno + context)
            for i in range(start, end + 1):
                result.append((i, lines[i-1].rstrip(), i == lineno))
            return result
        except Exception:
            return []

    def run(self) -> int:
        """Run the script and return exit code."""
        # Set up timeout
        if self.timeout > 0:
            def _timeout_handler():
                print(f"\n[pybreakz] TIMEOUT: Script exceeded {self.timeout}s", file=sys.stderr)
                os.kill(os.getpid(), signal.SIGTERM)
            timer = threading.Timer(self.timeout, _timeout_handler)
            timer.daemon = True
            timer.start()

        # Prepare script environment
        script_globals = {
            "__name__": "__main__",
            "__file__": self.script_path,
            "__builtins__": __builtins__,
        }

        exit_code = 0
        sys.settrace(self.trace_calls)

        try:
            with open(self.script_path) as f:
                source = f.read()
            code = compile(source, self.script_path, "exec")
            exec(code, script_globals)

        except SystemExit as e:
            exit_code = e.code if isinstance(e.code, int) else 0

        except Exception as e:
            exit_code = 1
            if self.on_exception:
                tb = sys.exc_info()[2]
                # Walk to innermost frame
                while tb.tb_next:
                    tb = tb.tb_next
                frame = tb.tb_frame
                exc_filepath = os.path.abspath(frame.f_code.co_filename)
                # Filter pybreakz frames from traceback text
                raw_tb = traceback.format_exc()
                _self = os.path.abspath(__file__)
                filtered_lines = []
                skip_next = False
                for line in raw_tb.splitlines():
                    if _self in line:
                        skip_next = True
                        continue
                    if skip_next and line.startswith("    "):
                        skip_next = False
                        continue
                    skip_next = False
                    filtered_lines.append(line)
                self.exception_info = {
                    "type": type(e).__name__,
                    "message": str(e),
                    "file": os.path.basename(frame.f_code.co_filename),
                    "line": tb.tb_lineno,
                    "traceback": "\n".join(filtered_lines),
                    "locals": fmt_locals(dict(frame.f_locals), self.watch, self.eval_exprs, frame),
                    "stack": fmt_frame(frame),
                    "source_line": self._get_source_line(tb.tb_lineno, exc_filepath),
                }
            else:
                # Still print the traceback for the user
                traceback.print_exc()

        finally:
            sys.settrace(None)
            if self.timeout > 0:
                timer.cancel()

        return exit_code

    def print_report(self):
        """Print the debug report to stdout."""
        if self.output_format == "json":
            self._print_json()
        else:
            self._print_text()

    def _print_text(self):
        script_name = os.path.basename(self.script_path)
        print(f"\n{'═' * 60}")
        print(f"  pybreakz report  →  {script_name}")
        print(f"{'═' * 60}")

        if not self.hits and not self.exception_info:
            print("\n  No breakpoints hit and no exception caught.")
            print("  (Check your line numbers — they must be executable lines)\n")
            return

        for i, hit in enumerate(self.hits, 1):
            if hit["type"] == "breakpoint":
                cond_str = f"  [condition: {hit['condition']}]" if hit.get("condition") else ""
                print(f"\n┌─ BREAKPOINT HIT #{i}  {hit['file']}:{hit['line']}{cond_str}")
                print(f"│  Source:  {hit['source_line']}")
                print("│")
                print("│  Call Stack (innermost last):")
                for s in hit["stack"]:
                    print(f"│{s}")
                print("│")
                for line in hit["locals"]:
                    print(f"│  {line}")
                print(f"└{SEPARATOR}")

            elif hit["type"] == "watchpoint":
                print(f"\n┌─ WATCHPOINT #{i}  {hit['var']}  changed  [{hit['file']}:{hit['line']}]")
                print(f"│  Source:  {hit['source_line']}")
                print(f"│  {hit['var']}: {hit['old']}  →  {hit['new']}")
                print("│")
                print("│  Call Stack (innermost last):")
                for s in hit["stack"]:
                    print(f"│{s}")
                print(f"└{SEPARATOR}")

        # Exception info
        if self.exception_info:
            ei = self.exception_info
            print(f"\n┌─ EXCEPTION  {ei['type']}: {ei['message']}")
            print(f"│  Location:  {ei['file']}:{ei['line']}")
            print(f"│  Source:    {ei['source_line']}")
            print("│")
            print("│  Call Stack:")
            for s in ei["stack"]:
                print(f"│{s}")
            print("│")
            for line in ei["locals"]:
                print(f"│  {line}")
            print("│")
            print("│  Full Traceback:")
            for line in ei["traceback"].splitlines():
                print(f"│  {line}")
            print(f"└{SEPARATOR}")

        print()

    def _print_json(self):
        output = {
            "script": self.script_path,
            "hits": self.hits,
            "exception": self.exception_info,
        }
        print(json.dumps(output, indent=2, default=str))


# ─── CLI ─────────────────────────────────────────────────────────────────────

def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="pybreakz",
        description="Debug Python scripts with breakpoints. Designed for Claude Code.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=textwrap.dedent("""
        Examples:
          pybreakz run script.py --breakpoints 42
          pybreakz run script.py --breakpoints 10,25,42 --watch "user_id,result"
          pybreakz run script.py --breakpoints "42:x>0,67:name=='foo'"
          pybreakz run script.py --on-exception --watch "self,request"
          pybreakz run script.py --breakpoints 10 --eval "len(items),type(r)"
          pybreakz run script.py --breakpoints 42 --format json
          pybreakz run script.py --breakpoints 10 -- --my-script-arg value

          # Cross-file breakpoints (set breakpoints in imported modules):
          pybreakz run script.py --breakpoints "utils.py:55,helpers.py:12"
          pybreakz run script.py --breakpoints "utils.py:55:x>0,42"

          # Variable watchpoints (record every time a variable changes):
          pybreakz run script.py --trace-vars "counter,result"
          pybreakz run script.py --trace-vars "x" --breakpoints 10
        """),
    )
    subparsers = parser.add_subparsers(dest="command")

    run_p = subparsers.add_parser("run", help="Run a script with debug tracing")
    run_p.add_argument("script", help="Python script to run")
    run_p.add_argument(
        "--breakpoints", "-b",
        help=(
            'Breakpoints to set. Formats: "LINE", "LINE:CONDITION", '
            '"FILE:LINE", "FILE:LINE:CONDITION". Comma-separated. '
            'e.g. "10,25" or "42:x>0" or "utils.py:55,helpers.py:12:done"'
        ),
    )
    run_p.add_argument(
        "--trace-vars", "-V",
        dest="trace_vars",
        help=(
            'Comma-separated variable names to watch. Records a hit whenever '
            'any of these variables change value, in any file. '
            'e.g. "counter,result,self.state"'
        ),
    )
    run_p.add_argument(
        "--watch", "-w",
        help='Comma-separated expressions to evaluate at each breakpoint. e.g. "x,len(items)"',
    )
    run_p.add_argument(
        "--eval", "-e", dest="eval_exprs",
        help='Extra expressions to evaluate at breakpoints. e.g. "df.shape,type(result)"',
    )
    run_p.add_argument(
        "--on-exception", "-x",
        action="store_true",
        help="Capture locals and stack on unhandled exception",
    )
    run_p.add_argument(
        "--timeout", "-t",
        type=int, default=30,
        help="Timeout in seconds (default: 30, 0 = no timeout)",
    )
    run_p.add_argument(
        "--format", "-f", dest="output_format",
        choices=["text", "json"], default="text",
        help="Output format (default: text)",
    )
    run_p.add_argument(
        "--max-hits",
        type=int, default=20,
        help="Max total hits to record across breakpoints and watchpoints (default: 20)",
    )
    return parser, run_p


def main():
    parser, run_p = build_parser()
    # We use parse_known_args so that script passthrough args (after --) don't
    # interfere with pybreakz's own flags.
    args, extra = parser.parse_known_args()

    if args.command is None:
        parser.print_help()
        sys.exit(0)

    if args.command == "run":
        # Validate script exists
        if not os.path.exists(args.script):
            print(f"[pybreakz] Error: Script not found: {args.script}", file=sys.stderr)
            sys.exit(1)

        script_basename = os.path.basename(args.script)

        # Parse breakpoints
        breakpoints: dict[str, dict[int, Optional[str]]] = {}
        if args.breakpoints:
            try:
                breakpoints = parse_breakpoints(args.breakpoints, script_basename)
            except ValueError as e:
                print(f"[pybreakz] Error parsing breakpoints: {e}", file=sys.stderr)
                sys.exit(1)

        # Parse trace-vars
        trace_vars = [v.strip() for v in args.trace_vars.split(",")] if args.trace_vars else []

        if not breakpoints and not args.on_exception and not trace_vars:
            print("[pybreakz] Warning: No breakpoints, --trace-vars, or --on-exception set.", file=sys.stderr)
            print("[pybreakz] Use --breakpoints LINE, --trace-vars NAME, or --on-exception", file=sys.stderr)

        # Parse watch/eval
        watch = [w.strip() for w in args.watch.split(",")] if args.watch else []
        eval_exprs = [e.strip() for e in args.eval_exprs.split(",")] if args.eval_exprs else []

        # Inject script args into sys.argv (extra = anything after --)
        script_args = extra
        if script_args and script_args[0] == "--":
            script_args = script_args[1:]
        sys.argv = [args.script] + script_args

        # Also add script dir to path so local imports work
        script_dir = os.path.dirname(os.path.abspath(args.script))
        if script_dir not in sys.path:
            sys.path.insert(0, script_dir)

        debugger = Debugger(
            script_path=args.script,
            breakpoints=breakpoints,
            watch=watch,
            eval_exprs=eval_exprs,
            on_exception=args.on_exception,
            timeout=args.timeout,
            output_format=args.output_format,
            max_hits=args.max_hits,
            trace_vars=trace_vars,
        )

        exit_code = debugger.run()
        debugger.print_report()
        sys.exit(exit_code)


if __name__ == "__main__":
    main()
