import io
import re
import signal
import sys
import threading
from typing import Any

from IPython.core.interactiveshell import InteractiveShell


MAX_STDOUT_LENGTH = 10_000
MAX_STDERR_LENGTH = 5_000


class PythonInstance:
    def __init__(self, session_id: str) -> None:
        self.session_id = session_id
        self.is_running = True
        self._execution_lock = threading.Lock()

        import os

        # Use configurable workspace path with fallback
        workspace = os.getenv("EXAAI_WORKSPACE", "/workspace")
        if os.path.isdir(workspace):
            os.chdir(workspace)

        # Add current directory and workspace to sys.path to ensure local imports work
        sys.path.insert(0, os.getcwd())
        if workspace not in sys.path:
            sys.path.insert(0, workspace)

        self.shell = InteractiveShell()
        self.shell.init_completer()
        self.shell.init_history()
        self.shell.init_logger()

        self._setup_proxy_functions()

    def _setup_proxy_functions(self) -> None:
        try:
            from exaaiagnt.tools.proxy import proxy_actions

            proxy_functions = [
                "list_requests",
                "list_sitemap",
                "repeat_request",
                "scope_rules",
                "send_request",
                "view_request",
                "view_sitemap_entry",
            ]

            proxy_dict = {name: getattr(proxy_actions, name) for name in proxy_functions}
            self.shell.user_ns.update(proxy_dict)
        except ImportError:
            # Proxy functions are optional, only available in full runtime
            pass

    def _validate_session(self) -> dict[str, Any] | None:
        if not self.is_running:
            return {
                "session_id": self.session_id,
                "stdout": "",
                "stderr": "Session is not running",
                "result": None,
            }
        return None

    def _setup_execution_environment(self, timeout: int) -> tuple[Any, io.StringIO, io.StringIO]:
        stdout_capture = io.StringIO()
        stderr_capture = io.StringIO()

        def timeout_handler(signum: int, frame: Any) -> None:
            raise TimeoutError(f"Code execution timed out after {timeout} seconds")

        old_handler = signal.signal(signal.SIGALRM, timeout_handler)
        signal.alarm(timeout)

        sys.stdout = stdout_capture
        sys.stderr = stderr_capture

        return old_handler, stdout_capture, stderr_capture

    def _cleanup_execution_environment(
        self, old_handler: Any, old_stdout: Any, old_stderr: Any
    ) -> None:
        signal.signal(signal.SIGALRM, old_handler)
        sys.stdout = old_stdout
        sys.stderr = old_stderr

    def _truncate_output(self, content: str, max_length: int, suffix: str) -> str:
        if len(content) > max_length:
            return content[:max_length] + suffix
        return content

    def _format_execution_result(
        self, execution_result: Any, stdout_content: str, stderr_content: str
    ) -> dict[str, Any]:
        # Strip ANSI codes from output
        ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
        stdout_content = ansi_escape.sub("", stdout_content)
        stderr_content = ansi_escape.sub("", stderr_content)

        stdout = self._truncate_output(
            stdout_content, MAX_STDOUT_LENGTH, "... [stdout truncated at 10k chars]"
        )

        if execution_result.result is not None:
            if stdout and not stdout.endswith("\n"):
                stdout += "\n"
            result_repr = repr(execution_result.result)
            result_repr = self._truncate_output(
                result_repr, MAX_STDOUT_LENGTH, "... [result truncated at 10k chars]"
            )
            stdout += result_repr

        stdout = self._truncate_output(
            stdout, MAX_STDOUT_LENGTH, "... [output truncated at 10k chars]"
        )

        stderr_content = stderr_content if stderr_content else ""
        stderr_content = self._truncate_output(
            stderr_content, MAX_STDERR_LENGTH, "... [stderr truncated at 5k chars]"
        )

        if (
            execution_result.error_before_exec or execution_result.error_in_exec
        ) and not stderr_content:
            stderr_content = "Execution error occurred"

        return {
            "session_id": self.session_id,
            "stdout": stdout,
            "stderr": stderr_content,
            "result": repr(execution_result.result)
            if execution_result.result is not None
            else None,
        }

    def _handle_execution_error(self, error: BaseException) -> dict[str, Any]:
        error_msg = str(error)
        error_msg = self._truncate_output(
            error_msg, MAX_STDERR_LENGTH, "... [error truncated at 5k chars]"
        )

        return {
            "session_id": self.session_id,
            "stdout": "",
            "stderr": error_msg,
            "result": None,
        }

    def execute_code(self, code: str, timeout: int = 30) -> dict[str, Any]:
        session_error = self._validate_session()
        if session_error:
            return session_error

        with self._execution_lock:
            old_stdout, old_stderr = sys.stdout, sys.stderr

            try:
                old_handler, stdout_capture, stderr_capture = self._setup_execution_environment(
                    timeout
                )

                try:
                    execution_result = self.shell.run_cell(code, silent=False, store_history=True)
                    signal.alarm(0)

                    return self._format_execution_result(
                        execution_result, stdout_capture.getvalue(), stderr_capture.getvalue()
                    )

                except (TimeoutError, KeyboardInterrupt, SystemExit) as e:
                    signal.alarm(0)
                    return self._handle_execution_error(e)

            finally:
                self._cleanup_execution_environment(old_handler, old_stdout, old_stderr)

    def close(self) -> None:
        self.is_running = False
        self.shell.reset(new_session=False)

    def is_alive(self) -> bool:
        return self.is_running
