"""
subprocess_runner — non-blocking subprocess primitives for the MCP sandbox.

All three entrypoints return the same dict shape as the legacy _run_inprocess():
    {
        "error":    bool,
        "stdout":   str,
        "stderr":   str,
        "images":   list,        # always [] for non-research invocations
        "_timeout": bool         # present and True only on timeout
    }

Callers can therefore pattern-match on {"_timeout", "error", "stdout", "stderr"}
uniformly regardless of whether the work ran in a subprocess or a thread.
"""

import asyncio
import json
import traceback
from pathlib import Path
from typing import Any, Callable


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _normalise(data: dict, stderr_fallback: str = "") -> dict:
    """Ensure the standard shape keys are present in a harness result dict."""
    data.setdefault("error", False)
    data.setdefault("stdout", "")
    data.setdefault("stderr", stderr_fallback)
    data.setdefault("images", [])
    return data


def _err_dict(stderr: str = "", stdout: str = "") -> dict:
    return {"error": True, "stdout": stdout, "stderr": stderr, "images": []}


def _timeout_dict() -> dict:
    return {"_timeout": True, "error": True, "stdout": "", "stderr": "", "images": []}


# ---------------------------------------------------------------------------
# Primitive 1: run_subprocess_argv
#
# Non-blocking equivalent of:
#   subprocess.run([sys.executable, harness, arg1, arg2, ...],
#                  capture_output=True, text=True, timeout=N, cwd=cwd)
#
# Used by: _execute_strategy, _run_research_harness
# ---------------------------------------------------------------------------

async def run_subprocess_argv(
    *cmd: str,
    timeout: int,
    cwd: Path | None = None,
) -> dict:
    """
    Spawn cmd as a subprocess, await completion, and return a normalised result dict.

    stdout is expected to contain a JSON object written by the harness.  It is
    decoded and normalised to the standard shape.  On JSON decode failure the
    raw stdout text is preserved in "stdout" and error is set to True.
    """
    try:
        proc = await asyncio.create_subprocess_exec(
            *cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            cwd=str(cwd) if cwd else None,
        )
        stdout_bytes, stderr_bytes = await asyncio.wait_for(
            proc.communicate(), timeout=timeout
        )
    except asyncio.TimeoutError:
        return _timeout_dict()
    except Exception as exc:
        return _err_dict(stderr=f"Harness launch failed: {exc}")

    stdout_text = stdout_bytes.decode(errors="replace")
    stderr_text = stderr_bytes.decode(errors="replace")

    if proc.returncode != 0:
        return _err_dict(
            stderr=f"Harness exited {proc.returncode}:\n{stderr_text}",
            stdout=stdout_text,
        )

    try:
        data = json.loads(stdout_text)
        return _normalise(data, stderr_fallback=stderr_text)
    except json.JSONDecodeError:
        return {"error": True, "stdout": stdout_text, "stderr": stderr_text, "images": []}


# ---------------------------------------------------------------------------
# Primitive 2: run_subprocess_stdin
#
# Non-blocking equivalent of the backtest pattern — JSON config fed via stdin.
# ---------------------------------------------------------------------------

async def run_subprocess_stdin(
    *cmd: str,
    stdin_data: bytes,
    timeout: int,
) -> dict:
    """
    Spawn cmd, write stdin_data to its stdin, await completion.

    Returns the same normalised dict shape as run_subprocess_argv.
    """
    try:
        proc = await asyncio.create_subprocess_exec(
            *cmd,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        stdout_bytes, stderr_bytes = await asyncio.wait_for(
            proc.communicate(stdin_data), timeout=timeout
        )
    except asyncio.TimeoutError:
        return _timeout_dict()
    except Exception as exc:
        return _err_dict(stderr=f"Harness launch failed: {exc}")

    stdout_text = stdout_bytes.decode(errors="replace")
    stderr_text = stderr_bytes.decode(errors="replace")

    if proc.returncode != 0:
        return _err_dict(
            stderr=f"Harness exited {proc.returncode}:\n{stderr_text}",
            stdout=stdout_text,
        )

    try:
        data = json.loads(stdout_text)
        return _normalise(data, stderr_fallback=stderr_text)
    except json.JSONDecodeError:
        return {"error": True, "stdout": stdout_text, "stderr": stderr_text, "images": []}


# ---------------------------------------------------------------------------
# Primitive 3: run_in_thread
#
# Async wrapper around asyncio.to_thread so the event loop stays responsive
# while CPU-bound or blocking-IO callables run in a worker thread.
#
# Used by: _execute_indicator (in-process indicator harness)
# ---------------------------------------------------------------------------

async def run_in_thread(
    fn: Callable,
    *args: Any,
    timeout: int,
) -> dict:
    """
    Run fn(*args) in a thread pool worker and yield to the event loop while waiting.

    On timeout the thread is abandoned (daemon) and _timeout_dict() is returned.
    On MemoryError or unexpected exception a standard error dict is returned.
    The returned dict is normalised to the standard shape.
    """
    from dexorder.memory_guard import cleanup_memory

    try:
        result = await asyncio.wait_for(
            asyncio.to_thread(fn, *args),
            timeout=timeout,
        )
        return _normalise(result)
    except asyncio.TimeoutError:
        return _timeout_dict()
    except MemoryError:
        cleanup_memory()
        return _err_dict(
            stderr="Script exceeded memory limit. Try reducing the data range or batch size."
        )
    except Exception:
        return _err_dict(stderr=traceback.format_exc())
