"""Python script execution in isolated processes.

Migrated and modernised from the old a2-svc-remote-py-executor.
Each script runs in a separate ``multiprocessing.Process`` for true
isolation — full Python stdlib access, unrestricted I/O.

Architecture reference: Section 13 (Python Script Execution on Remote).
"""

from __future__ import annotations

import asyncio
import base64
import contextlib
import io
import json
import logging
import multiprocessing as mp
import time
import traceback
from datetime import datetime, timezone
from typing import Any

logger = logging.getLogger("qremotex.execution")

# ──────────────────────────────────────────────── constants ──
MAX_LOG_LINE_BYTES = 10 * 1024  # 10 KB per log line
POLL_SLEEP_SEC = 0.05


# ──────────────────────────────────────────────── worker (runs in child process) ──

def _exec_worker(
    request_id: str,
    code_text: str,
    log_q: mp.Queue,
    result_q: mp.Queue,
    input_variables: dict[str, Any] | None = None,
    output_variable: str = "result",
) -> None:
    """Execute a Python script in an isolated child process.

    This function runs in a separate OS process.  It captures stdout/stderr
    and forwards them as log entries through ``log_q``.  The final result is
    sent via ``result_q``.

    ``output_variable`` controls which variable name is read from the script
    scope after execution (default ``"result"``).  qRunX sends this via the
    QRP ``config.output_variable`` field so that user scripts can use either
    ``output = ...`` or ``result = ...`` depending on node configuration.
    """
    input_variables = input_variables or {}

    # -- Logging helpers --
    def _emit(level: str, message: str, stream: str = "stdout") -> None:
        if len(message) > MAX_LOG_LINE_BYTES:
            message = message[:MAX_LOG_LINE_BYTES] + "... [truncated]"
        entry = {
            "type": "log",
            "level": level,
            "stream": stream,
            "message": message,
            "request_id": request_id,
            "timestamp": datetime.now(timezone.utc).isoformat(),
        }
        try:
            log_q.put_nowait(entry)
        except Exception:
            pass

    class _LineBuffer(io.TextIOBase):
        def __init__(self, level: str, stream_name: str):
            self.buf = ""
            self.level = level
            self.stream_name = stream_name

        def write(self, s: str) -> int:
            if not isinstance(s, str):
                s = str(s)
            self.buf += s
            while "\n" in self.buf:
                line, self.buf = self.buf.split("\n", 1)
                if line:
                    _emit(self.level, line, self.stream_name)
            return len(s)

        def flush(self) -> None:
            if self.buf:
                _emit(self.level, self.buf, self.stream_name)
                self.buf = ""

    # SafeConsole for structured logging inside scripts
    class SafeConsole:
        @staticmethod
        def log(*args: Any) -> None:
            _emit("info", " ".join(str(a) for a in args), "stdout")

        @staticmethod
        def info(*args: Any) -> None:
            _emit("info", " ".join(str(a) for a in args), "stdout")

        @staticmethod
        def warn(*args: Any) -> None:
            _emit("warn", " ".join(str(a) for a in args), "stdout")

        @staticmethod
        def error(*args: Any) -> None:
            _emit("error", " ".join(str(a) for a in args), "stderr")

    stdout_cap = _LineBuffer("info", "stdout")
    stderr_cap = _LineBuffer("error", "stderr")

    # Build execution globals — inject input dict + spread as top-level
    glb: dict[str, Any] = {
        "__name__": "__main__",
        "__package__": None,
        "__doc__": None,
        "__file__": "<remote_script>",
        "input": input_variables,
        "console": SafeConsole(),
    }
    # Spread input variables as top-level globals for convenience
    for k, v in input_variables.items():
        if k not in glb:
            glb[k] = v

    loc: dict[str, Any] = {}

    try:
        with contextlib.redirect_stdout(stdout_cap), contextlib.redirect_stderr(stderr_cap):
            compiled = compile(code_text, "<remote_script>", "exec")
            exec(compiled, glb, loc)  # noqa: S102 — intentional unrestricted exec

        stdout_cap.flush()
        stderr_cap.flush()

        # Extract result from script scope using the same priority as
        # qRunX local execution: check result, _result, output in order.
        # The configured output_variable (e.g. "code_op") is the context
        # storage key, not the script variable name, so we don't rely on it.
        _candidates = ["result", "_result", "output"]
        result = None
        # First pass: find a non-None value
        for _vname in _candidates:
            for _scope in (loc, glb):
                if _vname in _scope and _scope[_vname] is not None:
                    result = _scope[_vname]
                    break
            if result is not None:
                break
        # Second pass: accept falsy-but-valid values (0, False, "", [])
        if result is None:
            for _vname in _candidates:
                for _scope in (loc, glb):
                    if _vname in _scope:
                        result = _scope[_vname]
                        break
                if result is not None:
                    break

        payload: dict[str, Any] = {"success": True, "request_id": request_id}
        try:
            json.dumps(result)
            payload["result"] = result
        except (TypeError, ValueError):
            payload["result"] = {"__repr__": repr(result)}

        result_q.put(payload)

    except Exception as e:
        stdout_cap.flush()
        stderr_cap.flush()
        tb = traceback.format_exc()
        result_q.put({
            "success": False,
            "request_id": request_id,
            "error": str(e),
            "trace": tb,
        })


# ──────────────────────────────────────────────── async wrapper ──

class ScriptExecutor:
    """Async wrapper around the multiprocessing script execution."""

    async def execute(
        self,
        request_id: str,
        script_b64: str,
        timeout_sec: int,
        input_vars: dict[str, Any] | None = None,
        *,
        on_log: Any = None,  # async callback(log_entry)
        output_variable: str = "result",
    ) -> dict[str, Any]:
        """Execute a base64-encoded Python script.

        Returns a dict with ``success``, ``result``/``error``, ``duration_ms``.
        """
        # Decode script
        try:
            code_text = base64.b64decode(script_b64).decode("utf-8")
        except Exception as e:
            return {
                "success": False,
                "request_id": request_id,
                "error": f"base64 decode error: {e}",
                "duration_ms": 0,
            }

        log_q: mp.Queue = mp.Queue()
        result_q: mp.Queue = mp.Queue()

        proc = mp.Process(
            target=_exec_worker,
            args=(request_id, code_text, log_q, result_q, input_vars, output_variable),
            daemon=True,
        )
        proc.start()
        start = time.time()

        collected_logs: list[str] = []

        try:
            while True:
                # Drain log queue
                while not log_q.empty():
                    entry = log_q.get_nowait()
                    # Collect log messages for console_logs
                    msg_text = entry.get("message", "")
                    if msg_text:
                        collected_logs.append(msg_text)
                    if on_log:
                        await on_log(entry)

                # Check for result
                if not result_q.empty():
                    res = result_q.get()
                    res["duration_ms"] = int((time.time() - start) * 1000)
                    res["console_logs"] = collected_logs
                    return res

                # Timeout check
                if time.time() - start > timeout_sec:
                    if proc.is_alive():
                        proc.terminate()
                        proc.join(timeout=2)
                        if proc.is_alive():
                            proc.kill()
                    return {
                        "success": False,
                        "request_id": request_id,
                        "error": f"Script execution timed out after {timeout_sec}s",
                        "duration_ms": int((time.time() - start) * 1000),
                    }

                await asyncio.sleep(POLL_SLEEP_SEC)
        finally:
            if proc.is_alive():
                try:
                    proc.terminate()
                    proc.join(timeout=1)
                except Exception:
                    pass
