#!/usr/bin/env python3
'''
annodex-kernel-exec - Execute code on a Jupyter kernel and return rich mime outputs.

Usage:
  python3 annodex-kernel-exec.py <connection_file> <code>

The connection_file is a Jupyter kernel connection JSON file (e.g.
~/.local/share/jupyter/runtime/kernel-{id}.json).

Output: JSON object with:
  { "outputs": [{ "output_type", "data": {...}, "text", "ename", ... }] }

This bridges the gap between the annodex Node.js API and the Jupyter kernel,
allowing the Codex agent to get rich outputs (PNG images, HTML tables, etc.)
from the same kernel session that powers the thebe notebook UI.
'''

import json
import sys
import os
import base64
import traceback

from jupyter_client import BlockingKernelClient
from jupyter_client.connect import find_connection_file
from jupyter_client.session import Session


def normalize_output(msg_type, content):
    '''Convert Jupyter kernel messages into a normalized output dict.'''
    out = {"output_type": msg_type}

    if msg_type == "stream":
        out["text"] = content.get("text", "")
        out["name"] = content.get("name", "stdout")
        return out

    if msg_type == "error":
        out["ename"] = content.get("ename", "")
        out["evalue"] = content.get("evalue", "")
        out["traceback"] = content.get("traceback", [])
        return out

    if msg_type in ("execute_result", "display_data", "update_display_data"):
        data = content.get("data", {})
        out["data"] = {}
        for mime, value in data.items():
            if isinstance(value, (str, bytes)):
                out["data"][mime] = value
            elif isinstance(value, list):
                out["data"][mime] = "".join(str(v) for v in value)
            else:
                out["data"][mime] = str(value)
        if "text/plain" in data:
            out["text"] = data["text/plain"] if isinstance(data["text/plain"], str) else "".join(data["text/plain"])
        out["execution_count"] = content.get("execution_count")
        out["metadata"] = content.get("metadata", {})
        return out

    out["data"] = {"text/plain": json.dumps(content, default=str)}
    return out


def execute(connection_file, code, timeout=600):
    '''Execute code on a Jupyter kernel and collect outputs.
    Default timeout: 600s (10 minutes) for large data operations.
    '''
    if not os.path.exists(connection_file):
        return {
            "error": "Connection file not found: " + connection_file,
            "outputs": [],
        }

    client = BlockingKernelClient()
    client.load_connection_file(connection_file)
    client.start_channels()

    try:
        client.wait_for_ready(timeout=5)
    except RuntimeError as e:
        return {
            "error": "Kernel not ready: " + str(e),
            "outputs": [],
        }

    outputs = []
    errors = []

    try:
        msg_id = client.execute(code, timeout=timeout)

        while True:
            try:
                msg = client.get_iopub_msg(timeout=timeout)
            except Exception as e:
                errors.append("IOPub timeout: " + str(e))
                break

            parent_id = msg.get("parent_header", {}).get("msg_id", "")
            if parent_id != msg_id:
                continue

            msg_type = msg.get("msg_type", "")
            content = msg.get("content", {})

            if msg_type == "status" and content.get("execution_state") == "idle":
                break

            if msg_type in ("stream", "display_data", "execute_result",
                            "update_display_data", "error"):
                outputs.append(normalize_output(msg_type, content))

    except Exception as e:
        errors.append("Execution error: " + str(e))
    finally:
        try:
            client.stop_channels()
        except Exception:
            pass

    result = {"outputs": outputs}
    if errors:
        result["errors"] = errors
    return result


def main():
    if len(sys.argv) < 3:
        print(json.dumps({
            "error": "Usage: annodex-kernel-exec.py <connection_file> <code>  (use '-' as code to read from stdin)",
            "outputs": [],
        }))
        sys.exit(1)

    connection_file = sys.argv[1]
    code_arg = sys.argv[2]

    if code_arg == "-":
        code = sys.stdin.read()
    else:
        code = code_arg

    result = execute(connection_file, code)
    print(json.dumps(result))


if __name__ == "__main__":
    main()