"""
ColabPilot Bridge — registers RPC callbacks with Google Colab's output system.

The Chrome extension calls these callbacks via google.colab.kernel.invokeFunction().
"""

import sys
import traceback
import contextlib
from io import StringIO

import IPython.display

# Shared namespace persists variables across executions
_exec_namespace = {"__name__": "__colabpilot__"}


def _colabpilot_input(prompt=""):
    """Replacement for input() — not available in remote execution."""
    raise RuntimeError(
        f"input('{prompt}') is not supported via ColabPilot remote execution. "
        "Replace input() calls with hardcoded values for testing."
    )


_exec_namespace["input"] = _colabpilot_input


def ping():
    """Heartbeat callback — confirms the bridge is alive."""
    return IPython.display.JSON({"status": "ready"})


def execute(code, language="python"):
    """
    Execute code and return captured stdout, stderr, and any error.
    Cell creation is handled by the Chrome extension (colab-cells.js).
    """
    # Cell creation is handled by the extension side.
    # set_next_input() does NOT work from RPC callbacks — no active cell context.

    # Handle IPython ! prefix shell commands (e.g., !nvidia-smi, !pip install)
    if language == "python" and code.lstrip().startswith("!"):
        language = "bash"
        code = code.lstrip()[1:]

    stdout_buf = StringIO()
    stderr_buf = StringIO()
    error_msg = None

    try:
        if language == "bash":
            import subprocess
            proc = subprocess.run(
                code,
                shell=True,
                capture_output=True,
                text=True,
                timeout=300,
            )
            stdout_buf.write(proc.stdout)
            stderr_buf.write(proc.stderr)
            if proc.returncode != 0:
                error_msg = f"Process exited with code {proc.returncode}"
                if proc.stderr:
                    error_msg += f"\n{proc.stderr}"
        else:
            compiled = compile(code, "<colabpilot>", "exec")
            with contextlib.redirect_stdout(stdout_buf), \
                 contextlib.redirect_stderr(stderr_buf):
                exec(compiled, _exec_namespace)
    except Exception:
        error_msg = traceback.format_exc()

    return IPython.display.JSON({
        "stdout": stdout_buf.getvalue(),
        "stderr": stderr_buf.getvalue(),
        "error": error_msg,
    })


def connect():
    """
    Register the ColabPilot bridge with Google Colab's RPC system
    and set a DOM marker so the Chrome extension can detect readiness.

    Call this once in any Colab cell:

        import colabpilot
        colabpilot.connect()
    """
    from google.colab import output  # type: ignore

    # Register RPC callbacks
    output.register_callback("colabpilot_ping", ping)
    output.register_callback("colabpilot_execute", execute)

    print("ColabPilot bridge ready")
    print(f"Python {sys.version}")

    # Insert a DOM marker so the extension knows this output iframe has
    # a configured port and registered callbacks.
    try:
        output.eval_js("""
            var marker = document.getElementById("colabpilot-ready");
            if (!marker) {
                marker = document.createElement("div");
                marker.id = "colabpilot-ready";
                marker.style.display = "none";
                document.body.appendChild(marker);
            }
            "ok";
        """)
        print("Extension bridge marker set.")
    except Exception as e:
        print(f"Note: eval_js marker failed ({e}), bridge may need manual reload.")
