"""Persistent Python kernel for ML-Dash Studio.

Reads JSON commands from stdin, executes code in a persistent namespace,
and writes JSON responses to stdout with a magic delimiter prefix.

Protocol:
  stdin  <- JSON commands (one per line)
  stdout -> Lines prefixed with DELIMITER are JSON messages
            Lines without prefix are raw output (forwarded as stdout)

Message types:
  -> kernel_ready              Kernel is ready for commands
  -> stdout {data, cell_id}    Captured print output
  -> stderr {data, cell_id}    Captured error output
  -> image {data, format, cell_id}  Base64-encoded plot image
  -> html {data, cell_id}      HTML output (DataFrame tables)
  -> done {cell_id, exit_code, duration_ms}  Execution finished
"""

import ast
import base64
import io
import json
import os
import signal
import sys
import traceback
import time

# Force non-interactive matplotlib backend before any user imports
os.environ["MPLBACKEND"] = "Agg"

DELIMITER = "\x00MLDASH_MSG\x00"

# Persistent namespace for exec() — variables survive between cells
_namespace = {"__builtins__": __builtins__}


def _send(msg):
    """Send a JSON message to the bridge via stdout with magic delimiter."""
    sys.__stdout__.write(DELIMITER + json.dumps(msg) + "\n")
    sys.__stdout__.flush()


class _OutputCapture:
    """Captures writes to stdout/stderr and sends them as JSON messages."""

    def __init__(self, stream_type, cell_id, execution_id):
        self.stream_type = stream_type
        self.cell_id = cell_id
        self.execution_id = execution_id

    def write(self, text):
        if text:
            _send({
                "type": self.stream_type,
                "data": text,
                "cell_id": self.cell_id,
                "execution_id": self.execution_id,
            })

    def flush(self):
        pass

    def isatty(self):
        return False


def _capture_plots(cell_id, execution_id):
    """Capture any open matplotlib figures as base64 PNG images."""
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        return

    for fig_num in plt.get_fignums():
        fig = plt.figure(fig_num)
        buf = io.BytesIO()
        fig.savefig(buf, format="png", bbox_inches="tight", dpi=100)
        buf.seek(0)
        data = base64.b64encode(buf.read()).decode("ascii")
        _send({
            "type": "image",
            "data": data,
            "format": "png",
            "cell_id": cell_id,
            "execution_id": execution_id,
        })
        buf.close()

    plt.close("all")


def _execute(code, cell_id, execution_id):
    """Execute code in the persistent namespace, capturing all output."""
    old_stdout = sys.stdout
    old_stderr = sys.stderr
    sys.stdout = _OutputCapture("stdout", cell_id, execution_id)
    sys.stderr = _OutputCapture("stderr", cell_id, execution_id)

    exit_code = 0
    try:
        tree = ast.parse(code)

        if tree.body and isinstance(tree.body[-1], ast.Expr):
            # Execute all statements except the last expression
            if len(tree.body) > 1:
                module = ast.Module(body=tree.body[:-1], type_ignores=[])
                exec(compile(module, "<cell>", "exec"), _namespace)

            # Eval the last expression and display its result
            last_expr = tree.body[-1]
            expr_code = compile(ast.Expression(last_expr.value), "<cell>", "eval")
            result = eval(expr_code, _namespace)

            if result is not None:
                _namespace["_"] = result

                # Check for DataFrame/Series -> render as HTML table
                is_rich = False
                try:
                    import pandas as pd
                    if isinstance(result, (pd.DataFrame, pd.Series)):
                        html = result.to_html(max_rows=50, max_cols=20)
                        # Send HTML through the protocol (not captured stdout)
                        sys.stdout = old_stdout
                        sys.stderr = old_stderr
                        _send({
                            "type": "html",
                            "data": html,
                            "cell_id": cell_id,
                            "execution_id": execution_id,
                        })
                        sys.stdout = _OutputCapture("stdout", cell_id, execution_id)
                        sys.stderr = _OutputCapture("stderr", cell_id, execution_id)
                        is_rich = True
                except ImportError:
                    pass

                if not is_rich:
                    sys.stdout.write(repr(result) + "\n")
        else:
            # No trailing expression — just exec everything
            exec(compile(code, "<cell>", "exec"), _namespace)

    except KeyboardInterrupt:
        exit_code = 1
        sys.stderr.write("\nKeyboardInterrupt\n")
    except SystemExit as e:
        exit_code = e.code if isinstance(e.code, int) else 1
    except Exception:
        exit_code = 1
        sys.stderr.write(traceback.format_exc())
    finally:
        sys.stdout = old_stdout
        sys.stderr = old_stderr

    # Capture matplotlib plots rendered during execution
    _capture_plots(cell_id, execution_id)

    return exit_code


def _handle_interrupt(signum, frame):
    """SIGINT raises KeyboardInterrupt in the current exec()."""
    raise KeyboardInterrupt()


def main():
    """Main kernel loop — read JSON commands from stdin, execute, respond."""
    global _namespace

    signal.signal(signal.SIGINT, _handle_interrupt)

    _send({"type": "kernel_ready"})

    while True:
        try:
            line = sys.stdin.readline()
        except KeyboardInterrupt:
            # SIGINT between commands — ignore, keep running
            continue

        if not line:
            # stdin closed — bridge terminated
            break

        line = line.strip()
        if not line:
            continue

        try:
            cmd = json.loads(line)
        except json.JSONDecodeError:
            continue

        cmd_type = cmd.get("type")

        if cmd_type == "execute":
            code = cmd.get("code", "")
            cell_id = cmd.get("cell_id", "")
            execution_id = cmd.get("execution_id", "")

            start = time.time()
            exit_code = _execute(code, cell_id, execution_id)
            duration_ms = int((time.time() - start) * 1000)

            _send({
                "type": "done",
                "cell_id": cell_id,
                "execution_id": execution_id,
                "exit_code": exit_code,
                "duration_ms": duration_ms,
            })

        elif cmd_type == "restart":
            _namespace = {"__builtins__": __builtins__}
            # Clear matplotlib state if loaded
            try:
                import matplotlib.pyplot as plt
                plt.close("all")
            except ImportError:
                pass
            _send({"type": "kernel_ready"})

        elif cmd_type == "ping":
            _send({"type": "pong"})


if __name__ == "__main__":
    main()
