# tools/code_execution.py - Sandboxed Python code execution in an isolated venv.
from __future__ import annotations

import io
import json
import os
import subprocess
import sys
import tempfile
import time

from langchain_core.tools import tool

from tools._logging import _debug_log, emit_chat_event
from tools.live_view_visuals import emit_live_view_preview, image_to_preview_base64
from tools.shell import _safe_subprocess_env
from utils.path_safety import is_plugins_dir_write_blocked, get_user_output_dir, PLUGINS_DIR

_FILE_OUTPUT_DIR = get_user_output_dir()
_AGENT_VENV_DIR = os.path.join(os.path.expanduser("~"), ".mugi", "venv")
_PYTHON_TEXT_OUTPUT_NAME = "mugi_python_output.txt"
_PYTHON_TEXT_OUTPUT_MAX_CHARS = 200_000


def _get_agent_venv_python() -> str:
    """Return path to the agent venv's Python binary; create ~/.mugi/venv if missing."""
    if sys.platform == "win32":
        python_exe = os.path.join(_AGENT_VENV_DIR, "Scripts", "python.exe")
    else:
        python_exe = os.path.join(_AGENT_VENV_DIR, "bin", "python")
    if os.path.isfile(python_exe):
        return python_exe
    os.makedirs(os.path.dirname(_AGENT_VENV_DIR), exist_ok=True)
    try:
        subprocess.run(
            [sys.executable, "-m", "venv", _AGENT_VENV_DIR],
            check=True,
            capture_output=True,
            text=True,
            timeout=120,
        )
    except subprocess.CalledProcessError as e:
        err = (e.stderr or e.stdout or str(e)).strip()
        raise RuntimeError(f"Failed to create agent venv at {_AGENT_VENV_DIR}: {err}") from e
    except FileNotFoundError:
        raise RuntimeError("Cannot create venv: python -m venv not available.") from None
    if not os.path.isfile(python_exe):
        raise RuntimeError(f"Agent venv was created but {python_exe!r} not found.")
    _debug_log(f"[run_python_code] Created agent venv at {_AGENT_VENV_DIR}", "DEBUG")
    return python_exe


def _run_python_code_script(file_output_dir: str, user_code: str, evolution_mode: bool = False) -> str:
    """Build script content: bootstrap (FILE_OUTPUT_DIR + open wrapper for write restrictions) then user code."""
    path_repr = repr(file_output_dir)
    home_repr = repr(os.path.expanduser("~"))
    plugins_repr = repr(os.path.abspath(os.path.expanduser(PLUGINS_DIR)))
    evolution_repr = repr(evolution_mode)
    bootstrap = f"""import os
import sys
FILE_OUTPUT_DIR = {path_repr}
_HOME = {home_repr}
_PLUGINS_DIR = {plugins_repr}
_EVOLUTION_MODE = {evolution_repr}
_ALLOWED_WRITE_PREFIXES = [FILE_OUTPUT_DIR]
if _EVOLUTION_MODE:
    _ALLOWED_WRITE_PREFIXES.append(_PLUGINS_DIR)
_ALLOWED_WRITE_PREFIXES = tuple(os.path.realpath(os.path.expanduser(p)) for p in _ALLOWED_WRITE_PREFIXES)
_ALLOWED_READ_PREFIXES = tuple(
    os.path.realpath(os.path.expanduser(p))
    for p in (
        FILE_OUTPUT_DIR, os.path.join(_HOME, "Documents"), os.path.join(_HOME, "Downloads"),
        "/tmp", "/var/folders", os.path.join(_HOME, ".mugi", "venv"),
        os.path.join(_HOME, ".matplotlib"), os.path.join(_HOME, ".config"),
        sys.prefix, sys.exec_prefix
    )
)
_real_open = open
def _sandbox_open(path, mode="r", *a, **k):
    if isinstance(path, int):
        return _real_open(path, mode, *a, **k)
    path_str = os.fspath(path) if not isinstance(path, str) else path
    path_abs = os.path.realpath(os.path.normpath(os.path.abspath(os.path.expanduser(path_str))))
    if "w" in mode or "a" in mode:
        if not any(path_abs == p or path_abs.startswith(p + os.sep) for p in _ALLOWED_WRITE_PREFIXES):
            raise OSError("Write blocked: use the FILE_OUTPUT_DIR variable. Example: open(FILE_OUTPUT_DIR + '/name.txt', 'w'). FILE_OUTPUT_DIR = " + repr(FILE_OUTPUT_DIR))
    elif not any(path_abs.startswith(p) for p in _ALLOWED_READ_PREFIXES):
        raise OSError("Sandbox read not allowed for this path.")
    return _real_open(path, mode, *a, **k)
open = _sandbox_open
import builtins as _builtins
_builtins.open = _sandbox_open

import subprocess as _sb
for _fn in ("run", "Popen", "call", "check_call", "check_output"):
    setattr(_sb, _fn, (lambda n: lambda *a, **k: (_ for _ in ()).throw(OSError(n + " is blocked in sandbox")))("subprocess." + _fn))
import socket as _sk
_sk._real_socket = _sk.socket
class _BlockedSocket(_sk._real_socket):
    def __init__(self, *a, **k):
        raise OSError("socket.socket is blocked in sandbox")
_sk.socket = _BlockedSocket
_sk.create_connection = lambda *a, **k: (_ for _ in ()).throw(OSError("socket.create_connection is blocked in sandbox"))
import shutil as _sh
_sh.rmtree = lambda *a, **k: (_ for _ in ()).throw(OSError("shutil.rmtree is blocked in sandbox"))
_sh.move = lambda *a, **k: (_ for _ in ()).throw(OSError("shutil.move is blocked in sandbox"))
for _fn in ("system", "popen", "popen2", "popen3", "popen4"):
    if hasattr(os, _fn):
        setattr(os, _fn, (lambda n: lambda *a, **k: (_ for _ in ()).throw(OSError(n + " is blocked in sandbox")))("os." + _fn))
for _fn in ("execl", "execle", "execlp", "execlpe", "execv", "execve", "execvp", "execvpe",
            "spawnl", "spawnle", "spawnlp", "spawnlpe", "spawnv", "spawnve", "spawnvp", "spawnvpe"):
    if hasattr(os, _fn):
        setattr(os, _fn, (lambda n: lambda *a, **k: (_ for _ in ()).throw(OSError(n + " is blocked in sandbox")))("os." + _fn))

try:
    import matplotlib as _mpl
    _mpl.rcParams["savefig.dpi"] = 200
    _mpl.rcParams["figure.dpi"] = 200
except Exception:
    pass

"""
    return bootstrap + user_code


def _write_python_text_preview(
    file_output_dir: str,
    stdout: str,
    stderr: str,
    *,
    max_chars: int = _PYTHON_TEXT_OUTPUT_MAX_CHARS,
) -> tuple[str, bool] | None:
    """Write captured text output to the Results preview file.

    Returns (basename, truncated) when a preview file was written.
    """
    sections = []
    if stdout:
        sections.append(f"stdout:\n{stdout}")
    if stderr:
        sections.append(f"stderr:\n{stderr}")
    if not sections:
        return None
    text = "\n\n".join(sections).strip() + "\n"
    truncated = False
    if len(text) > max_chars:
        text = (
            text[:max_chars].rstrip()
            + f"\n\n[Output truncated to {max_chars} characters for preview.]\n"
        )
        truncated = True
    os.makedirs(file_output_dir, exist_ok=True)
    path = os.path.join(file_output_dir, _PYTHON_TEXT_OUTPUT_NAME)
    with open(path, "w", encoding="utf-8", errors="replace") as f:
        f.write(text)
    return _PYTHON_TEXT_OUTPUT_NAME, truncated


@tool
def run_python_code(code: str) -> str:
    """Execute a Python code snippet in an isolated venv. Use for math, data analysis, charts, diagrams, flowcharts, and text outputs. The venv already has numpy, pandas, matplotlib, and seaborn — do NOT install anything, just write code. FILE_OUTPUT_DIR is a pre-defined Python variable already set to ~/Documents/Mugi — use it directly (do NOT hardcode any path). Example: plt.savefig(FILE_OUTPUT_DIR + '/my_chart.png'). Charts are automatically displayed inline in the chat; text stdout is automatically shown in the Results window when no chart is produced. Do NOT tell the user a file path or ask them to open the file — just describe the result. For flowcharts and diagrams, use matplotlib patches and annotations. Input is the Python code string."""
    _debug_log(f"run_python_code: running snippet ({len(code) if code else 0} chars)", "INFO")
    _debug_log(f"[DEBUG run_python_code] Called with code length: {len(code) if code else 0}")
    if not code or not isinstance(code, str):
        return "Invalid input: 'code' must be a non-empty string."
    if "kraken" in code.lower():
        _debug_log("[run_python_code] Code mentions 'kraken'. Use ocr_with_kraken_api tool instead.", "WARN")
    try:
        venv_python = _get_agent_venv_python()
    except RuntimeError as e:
        _debug_log(f"[ERROR run_python_code] {e}", "ERROR")
        return f"Agent venv setup failed: {e}"
    _file_output_dir = get_user_output_dir()
    try:
        from core.config import evolution_enabled
        _evolution_on = evolution_enabled()
    except Exception:
        _evolution_on = False
    emit_live_view_preview("Running Python code…", tool_name="run_python_code")
    script_content = _run_python_code_script(_file_output_dir, code, evolution_mode=_evolution_on)
    fd, script_path = tempfile.mkstemp(suffix=".py", prefix="mugi_run_python_")
    try:
        try:
            os.write(fd, script_content.encode("utf-8"))
        finally:
            os.close(fd)
        env = _safe_subprocess_env(FILE_OUTPUT_DIR=_file_output_dir, MPLBACKEND="Agg")
        cwd = os.path.expanduser("~")
        if not os.path.isdir(cwd):
            cwd = _file_output_dir
        result = subprocess.run(
            [venv_python, script_path],
            env=env,
            cwd=cwd,
            capture_output=True,
            text=True,
            timeout=30,
        )
        output = (result.stdout or "").strip()
        error = (result.stderr or "").strip()
        candidates = []
        if result.returncode == 0 and os.path.isdir(_file_output_dir):
            _img_ext = (".png", ".jpg", ".jpeg")
            _cutoff = time.time() - 60
            try:
                for f in os.listdir(_file_output_dir):
                    if f.lower().endswith(_img_ext):
                        p = os.path.join(_file_output_dir, f)
                        if os.path.isfile(p) and os.path.getmtime(p) >= _cutoff:
                            candidates.append((os.path.getmtime(p), p))
                candidates.sort(key=lambda x: -x[0])
                if candidates:
                    with open(candidates[0][1], "rb") as _f:
                        b64 = image_to_preview_base64(_f.read(), max_width=500)
                    if b64:
                        emit_live_view_preview(
                            f"Output: {os.path.basename(candidates[0][1])}",
                            tool_name="run_python_code",
                            preview_image_base64=b64,
                        )
                    chart_filenames = [os.path.basename(c[1]) for c in candidates]
                    emit_chat_event({"type": "image_add", "files": chart_filenames})
                    emit_chat_event({
                        "type": "display_window",
                        "files": chart_filenames[:1],
                        "title": "Chart",
                        "content_type": "image",
                    })
            except Exception:
                pass
        text_preview = None
        if result.returncode == 0 and not candidates:
            try:
                text_preview = _write_python_text_preview(_file_output_dir, output, error)
                if text_preview:
                    emit_chat_event({
                        "type": "display_window",
                        "files": [text_preview[0]],
                        "title": "Python Output",
                        "content_type": "text",
                    })
            except Exception:
                _debug_log("[run_python_code] Failed to write text preview", "WARN")
        parts = []
        if output:
            parts.append(f"stdout:\n{output}\n")
        if error:
            parts.append(f"stderr:\n{error}\n")
        if result.returncode == 0 and candidates:
            saved_names = [os.path.basename(c[1]) for c in candidates]
            parts.append(
                f"Charts saved and ALREADY DISPLAYED INLINE in the chat: {', '.join(saved_names)}. "
                "Do NOT tell the user to open or find these files — they can already see them. "
                "Just describe what the chart shows."
            )
        elif text_preview:
            extra = " Output was truncated for preview." if text_preview[1] else ""
            parts.append(
                "Text output saved and ALREADY DISPLAYED in the Results window. "
                f"Do NOT tell the user to open or find the file.{extra} "
                "Just summarize the output."
            )
        parts.append(f"Return code: {result.returncode}")
        return "".join(parts) if parts else "Executed successfully (no output)."
    except subprocess.TimeoutExpired:
        return "Execution timed out after 30 seconds (possible infinite loop)."
    except Exception as e:
        _debug_log(f"[ERROR run_python_code] {e}", "ERROR")
        return f"Code execution failed: {str(e)}"
    finally:
        try:
            os.unlink(script_path)
        except OSError:
            pass
