#!/usr/bin/env python3
"""
Notebook execution runner for agent tasks.
Wraps papermill with sensible defaults and structured JSON reporting.

Usage as CLI:
    python -m devkit.notebook_runner notebooks/task.ipynb output.ipynb -p threshold 1.0
    python -m devkit.notebook_runner notebooks/task.ipynb --dry-run

Usage as module:
    from devkit.notebook_runner import run_notebook
    result = run_notebook("task.ipynb", "output.ipynb", parameters={"threshold": 1.0})
"""

import argparse
import json
import os
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Optional


def run_notebook(
    input_path: str,
    output_path: Optional[str] = None,
    parameters: Optional[Dict[str, Any]] = None,
    kernel: str = "python3",
    timeout: Optional[int] = None,
    dry_run: bool = False,
) -> Dict[str, Any]:
    input_path = str(Path(input_path).expanduser().resolve())
    if output_path is None:
        fd, output_path = tempfile.mkstemp(suffix=".ipynb")
        os.close(fd)
    else:
        output_path = str(Path(output_path).expanduser().resolve())

    if dry_run:
        return {
            "status": "dry_run",
            "input": input_path,
            "output": output_path,
            "parameters": parameters or {},
            "kernel": kernel,
        }

    cmd = ["papermill", input_path, output_path, "--kernel", kernel]
    if timeout is not None:
        cmd.extend(["--execution-timeout", str(timeout)])
    for key, value in (parameters or {}).items():
        cmd.extend(["-p", key, str(value)])

    start = os.times()
    try:
        proc = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            check=True,
        )
        end = os.times()
        return {
            "status": "ok",
            "input": input_path,
            "output": output_path,
            "parameters": parameters or {},
            "kernel": kernel,
            "elapsed_seconds": round((end.elapsed - start.elapsed), 2),
            "stdout": proc.stdout[-2000:] if proc.stdout else "",
        }
    except subprocess.CalledProcessError as e:
        end = os.times()
        return {
            "status": "error",
            "input": input_path,
            "output": output_path,
            "parameters": parameters or {},
            "kernel": kernel,
            "elapsed_seconds": round((end.elapsed - start.elapsed), 2),
            "stdout": e.stdout[-2000:] if e.stdout else "",
            "stderr": e.stderr[-2000:] if e.stderr else "",
            "returncode": e.returncode,
        }
    except FileNotFoundError:
        return {
            "status": "error",
            "message": "papermill not found. Install with: uv tool install papermill",
        }


def main(argv: List[str] = None) -> int:
    argv = argv or sys.argv[1:]
    parser = argparse.ArgumentParser(description="Notebook runner for agents")
    parser.add_argument("input", help="Input notebook path")
    parser.add_argument("output", nargs="?", default=None, help="Output notebook path")
    parser.add_argument("-p", "--parameter", action="append", default=[], help="Parameters as key=value")
    parser.add_argument("--kernel", default="python3")
    parser.add_argument("--timeout", type=int, default=None)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args(argv)

    parameters = {}
    for raw in args.parameter:
        if "=" not in raw:
            print(f"Invalid parameter (expected key=value): {raw}", file=sys.stderr)
            return 1
        k, v = raw.split("=", 1)
        # Best-effort type inference
        if v.lower() in ("true", "false"):
            v = v.lower() == "true"
        else:
            try:
                v = int(v)
            except ValueError:
                try:
                    v = float(v)
                except ValueError:
                    pass
        parameters[k] = v

    result = run_notebook(
        args.input,
        args.output,
        parameters=parameters,
        kernel=args.kernel,
        timeout=args.timeout,
        dry_run=args.dry_run,
    )
    print(json.dumps(result, indent=2))
    return 0 if result.get("status") == "ok" else 1


if __name__ == "__main__":
    sys.exit(main())