"""Sandboxed Python code execution using RestrictedPython.

Handles compilation and execution of LLM-generated code in a safe environment.
"""

from __future__ import annotations

import json
import os
import re
import traceback
from collections import Counter, defaultdict
from datetime import datetime, timedelta
from io import StringIO
from typing import Any, TYPE_CHECKING

from RestrictedPython import compile_restricted, safe_globals
from RestrictedPython.Eval import default_guarded_getiter, default_guarded_getitem
from RestrictedPython.Guards import guarded_iter_unpack_sequence, safer_getattr

from .exceptions import SandboxError
from .execution_context import DEFAULT_MAX_OUTPUT_CHARS
from .message_database import cf_to_datetime

if TYPE_CHECKING:
    from .message_database import MessageDatabase
    from .llm_orchestrator import LlmOrchestrator


class RestrictedExecutor:
    """Handles sandboxed Python code execution.

    Responsibilities:
    - Compile code with RestrictedPython
    - Build safe execution namespace
    - Execute code with captured stdout
    """

    # Safe builtins allowed in sandbox
    SAFE_BUILTINS: dict[str, Any] = {
        "True": True,
        "False": False,
        "None": None,
        "abs": abs,
        "all": all,
        "any": any,
        "bool": bool,
        "dict": dict,
        "enumerate": enumerate,
        "filter": filter,
        "float": float,
        "frozenset": frozenset,
        "int": int,
        "isinstance": isinstance,
        "len": len,
        "list": list,
        "map": map,
        "max": max,
        "min": min,
        "range": range,
        "reversed": reversed,
        "round": round,
        "set": set,
        "sorted": sorted,
        "str": str,
        "sum": sum,
        "tuple": tuple,
        "zip": zip,
        "hasattr": hasattr,
    }

    _DEBUG_ENV = "RECALL_RLM_DEBUG"
    _DEBUG_LOG_ENV = "RECALL_RLM_DEBUG_LOG_PATH"
    _DEBUG_LOG_DEFAULT = "/tmp/recall_rlm_sandbox_error.jsonl"
    _DEBUG_CODE_LIMIT = 20000

    def __init__(
        self,
        db: MessageDatabase,
        llm: LlmOrchestrator | None = None,
        answer: dict[str, Any] | None = None,
        max_output_chars: int = DEFAULT_MAX_OUTPUT_CHARS,
    ) -> None:
        self._db = db
        self._llm = llm
        # Answer dict for iterative refinement (Prime Intellect RLM pattern)
        # - answer["content"]: Accumulated answer text (read/write by sandbox code)
        # - answer["ready"]: Set to True when answer is complete
        self._answer = answer if answer is not None else {"content": "", "ready": False}
        self._max_output_chars = max(1, int(max_output_chars))

    def _maybe_log_debug(self, code: str, exc: Exception, trace: str | None = None) -> None:
        if not os.environ.get(self._DEBUG_ENV):
            return
        try:
            log_path = os.environ.get(self._DEBUG_LOG_ENV, self._DEBUG_LOG_DEFAULT)
            payload = {
                "error_type": exc.__class__.__name__,
                "message": str(exc),
                "traceback": trace or "",
                "code": code[: self._DEBUG_CODE_LIMIT],
            }
            with open(log_path, "a", encoding="utf-8") as handle:
                handle.write(json.dumps(payload, ensure_ascii=True))
                handle.write("\n")
        except Exception:
            return

    class _OutputBuffer:
        def __init__(self, max_chars: int) -> None:
            self._buffer = StringIO()
            self._max_chars = max_chars
            self._size = 0

        def write(self, text: str) -> int:
            if not text:
                return 0
            new_size = self._size + len(text)
            if new_size > self._max_chars:
                raise SandboxError(
                    f"Sandbox output exceeded {self._max_chars} characters (attempted {new_size}). "
                    "Summarize or aggregate results instead of printing raw rows."
                )
            self._size = new_size
            return self._buffer.write(text)

        def flush(self) -> None:
            self._buffer.flush()

        def getvalue(self) -> str:
            return self._buffer.getvalue()

    def _build_globals(self, output: StringIO, extra_globals: dict[str, Any] | None = None) -> dict[str, Any]:
        """Build the restricted globals namespace."""
        def safe_print(*args: Any, **kwargs: Any) -> None:
            kwargs["file"] = output
            print(*args, **kwargs)

        def guarded_write(obj: Any) -> Any:
            if isinstance(obj, (dict, list)):
                return obj
            raise SandboxError(f"Cannot assign items on {type(obj).__name__}")

        restricted_globals = safe_globals.copy()

        # RestrictedPython guards
        restricted_globals["_getiter_"] = default_guarded_getiter
        restricted_globals["_getitem_"] = default_guarded_getitem
        restricted_globals["_iter_unpack_sequence_"] = guarded_iter_unpack_sequence
        restricted_globals["_getattr_"] = safer_getattr
        restricted_globals["_write_"] = guarded_write

        # Safe builtins with print
        builtins = dict(self.SAFE_BUILTINS)
        builtins["print"] = safe_print
        builtins["getattr"] = safer_getattr
        restricted_globals["__builtins__"] = builtins

        # Safe modules
        restricted_globals["re"] = re
        restricted_globals["json"] = json
        restricted_globals["Counter"] = Counter
        restricted_globals["defaultdict"] = defaultdict
        restricted_globals["datetime"] = datetime
        restricted_globals["timedelta"] = timedelta
        restricted_globals["cf_to_datetime"] = cf_to_datetime

        # Database functions
        restricted_globals["query"] = self._db.query
        restricted_globals["get_messages"] = self._db.get_messages
        restricted_globals["get_stats"] = self._db.get_stats
        restricted_globals["get_chats"] = self._db.get_chats
        restricted_globals["search"] = self._db.search
        restricted_globals["get_recent"] = self._db.get_recent
        restricted_globals["get_chat_messages"] = self._db.get_chat_messages
        restricted_globals["sample"] = self._db.sample

        # LLM functions (if orchestrator provided)
        if self._llm:
            restricted_globals["llm_query"] = self._llm.llm_query
            restricted_globals["chunk_data"] = self._llm.chunk_data
            restricted_globals["batch_llm_query"] = self._llm.batch_llm_query

        # Answer dict for iterative refinement (Prime Intellect RLM pattern)
        # The sandbox code can modify answer["content"] and set answer["ready"] = True
        restricted_globals["answer"] = self._answer

        # Print handler for RestrictedPython
        class PrintHandler:
            def __init__(self, _getattr_: Any = None) -> None:
                self._getattr_ = _getattr_

            def _call_print(self, *args: Any, **kwargs: Any) -> None:
                text = " ".join(str(arg) for arg in args)
                end = kwargs.get("end", "\n")
                output.write(text + end)

            def __call__(self) -> str:
                return ""

        restricted_globals["_print_"] = PrintHandler

        if extra_globals:
            restricted_globals.update(extra_globals)

        return restricted_globals

    def close(self) -> None:
        """Close underlying database connection."""
        self._db.close()

    def run_code(self, code: str, extra_globals: dict[str, Any] | None = None) -> str:
        """Execute Python code in a restricted environment.

        Returns the captured stdout output.

        Raises:
            SandboxError: If compilation or execution fails
        """
        if "__import__" in code or "import " in code:
            raise SandboxError("Imports are not allowed in this environment. Use provided helpers instead.")

        if extra_globals and "answer" in extra_globals:
            self._answer = extra_globals["answer"]

        output = self._OutputBuffer(self._max_output_chars)
        restricted_globals = self._build_globals(output, extra_globals=extra_globals)

        # Compile with RestrictedPython
        try:
            compiled_code = compile_restricted(
                code,
                filename="<rlm>",
                mode="exec",
            )
        except SyntaxError as e:
            self._maybe_log_debug(code, e, traceback.format_exc())
            raise SandboxError(f"Syntax error in code: {e}")

        if hasattr(compiled_code, 'errors') and compiled_code.errors:
            error = SandboxError(f"Compilation errors: {compiled_code.errors}")
            self._maybe_log_debug(code, error)
            raise SandboxError(f"Compilation errors: {compiled_code.errors}")

        if hasattr(compiled_code, 'code'):
            compiled_code = compiled_code.code

        if compiled_code is None:
            error = SandboxError("Compilation failed: no code generated")
            self._maybe_log_debug(code, error)
            raise SandboxError("Compilation failed: no code generated")

        # Execute the RestrictedPython-compiled bytecode
        # This is safe: RestrictedPython blocks dangerous operations at compile time
        try:
            exec_locals: dict[str, Any] = {}
            # Get the real builtin executor from builtins (RestrictedPython usage pattern)
            builtins_module = __builtins__ if isinstance(__builtins__, dict) else vars(__builtins__)
            builtin_executor = builtins_module.get("exec")
            if builtin_executor is None:
                error = SandboxError("Could not find builtin executor")
                self._maybe_log_debug(code, error)
                raise error
            builtin_executor(compiled_code, restricted_globals, exec_locals)
        except SandboxError as exc:
            self._maybe_log_debug(code, exc, traceback.format_exc())
            raise
        except Exception as e:
            trace = traceback.format_exc()
            self._maybe_log_debug(code, e, trace)
            if os.environ.get("RECALL_RLM_DEBUG"):
                trace = traceback.format_exc()
                raise SandboxError(
                    f"Execution error: {type(e).__name__}: {e}\nTraceback:\n{trace}"
                )
            raise SandboxError(f"Execution error: {type(e).__name__}: {e}")

        return output.getvalue()

    # Alias for backward compatibility
    def execute(self, code: str, extra_globals: dict[str, Any] | None = None) -> str:
        """Alias for run_code() for backward compatibility."""
        return self.run_code(code, extra_globals=extra_globals)

    @property
    def answer(self) -> dict[str, Any]:
        """Access the answer dict for iterative refinement."""
        return self._answer

    def set_answer(self, answer: dict[str, Any]) -> None:
        """Update the answer dict (for resetting between sessions)."""
        self._answer = answer
