"""Python executor for environments without nsjail.

This executor uses process-level isolation and runs user scripts in a dedicated
PID namespace when available. It is intended for environments where nsjail
cannot run (for example, without SYS_ADMIN).
"""

import asyncio
import hashlib
import json
import logging
import os
import shutil
import subprocess
import tempfile
import time
from contextlib import suppress
from pathlib import Path
from typing import Any

from tracecat.config import (
    TRACECAT__SANDBOX_CACHE_DIR,
    TRACECAT__SANDBOX_DEFAULT_TIMEOUT,
    TRACECAT__SANDBOX_PYPI_EXTRA_INDEX_URLS,
    TRACECAT__SANDBOX_PYPI_INDEX_URL,
)
from tracecat.logger import logger
from tracecat.sandbox.exceptions import (
    PackageInstallError,
    SandboxExecutionError,
    SandboxTimeoutError,
)
from tracecat.sandbox.types import SandboxResult

module_logger = logging.getLogger(__name__)

SAFE_WRAPPER_SCRIPT = '''
import asyncio
import importlib
import inspect
import json
import os
import sys
import traceback
from pathlib import Path

def _install_action_gateway_sdk_transport():
    socket_path = os.environ.get("TRACECAT__ACTION_GATEWAY_SOCKET")
    if not socket_path:
        return

    try:
        import httpx

        sdk_client = importlib.import_module("tracecat_registry.sdk.client")
    except ImportError:
        return

    tracecat_client_cls = getattr(sdk_client, "TracecatClient", None)
    if tracecat_client_cls is None:
        return

    if hasattr(tracecat_client_cls, "_request_url_and_transport"):
        return

    if getattr(tracecat_client_cls, "_tracecat_action_gateway_transport", False):
        return

    async def request(
        self,
        method,
        path,
        *,
        params=None,
        json=None,
        headers=None,
    ):
        request_headers = self._get_headers()
        if headers:
            request_headers.update(headers)

        async with httpx.AsyncClient(
            transport=httpx.AsyncHTTPTransport(uds=socket_path),
            timeout=getattr(self, "_timeout", 120.0),
        ) as client:
            response = await client.request(
                method,
                "http://tracecat-action-gateway/internal" + path,
                params=params,
                json=json,
                headers=request_headers,
            )

        if not response.is_success:
            self._handle_error_response(response)

        if not response.content:
            return None

        return response.json()

    tracecat_client_cls.request = request
    tracecat_client_cls._tracecat_action_gateway_transport = True

def _init_tracecat_context():
    _install_action_gateway_sdk_transport()
    try:
        from tracecat_registry.context import init_context_from_env
    except ImportError:
        return
    try:
        init_context_from_env()
    except ValueError:
        return

def _resolve_output(value):
    if not inspect.isawaitable(value):
        return value

    async def await_value():
        return await value

    return asyncio.run(await_value())

def main():
    """Execute user script and capture results."""
    work_dir = "{work_dir}"

    # Read inputs from file
    inputs_path = Path(work_dir) / "inputs.json"
    if inputs_path.exists():
        inputs = json.loads(inputs_path.read_text())
    else:
        inputs = {{}}

    result = {{
        "success": False,
        "output": None,
        "error": None,
        "traceback": None,
        "stdout": "",
        "stderr": "",
    }}

    # Capture stdout/stderr
    import io
    old_stdout = sys.stdout
    old_stderr = sys.stderr
    sys.stdout = io.StringIO()
    sys.stderr = io.StringIO()

    try:
        # Read and execute the user script
        script_path = Path(work_dir) / "script.py"
        script_code = script_path.read_text()
        _init_tracecat_context()
        script_globals = {{"__name__": "__main__", "__file__": str(script_path)}}
        exec(script_code, script_globals)

        # Find the callable function
        main_func = script_globals.get("main")
        if main_func is None:
            for name, obj in script_globals.items():
                if inspect.isfunction(obj) and not name.startswith("_"):
                    main_func = obj
                    break

        if main_func is None:
            raise ValueError("No callable function found in script")

        # Call the function with inputs
        if inputs:
            call = main_func(**inputs)
        else:
            call = main_func()
        output = _resolve_output(call)

        result["success"] = True
        result["output"] = output

    except Exception as e:
        result["error"] = f"{{type(e).__name__}}: {{e}}"
        result["traceback"] = traceback.format_exc()

    finally:
        result["stdout"] = sys.stdout.getvalue()
        result["stderr"] = sys.stderr.getvalue()
        sys.stdout = old_stdout
        sys.stderr = old_stderr

    # Write result to file
    result_path = Path(work_dir) / "result.json"
    try:
        result_path.write_text(json.dumps(result))
    except (TypeError, ValueError):
        # Output not JSON-serializable, convert to repr
        result["output"] = repr(result["output"])
        try:
            result_path.write_text(json.dumps(result))
        except Exception as e:
            result["output"] = None
            result["error"] = f"Output not JSON-serializable: {{type(e).__name__}}: {{e}}"
            result["success"] = False
            result_path.write_text(json.dumps(result))

    sys.exit(0 if result["success"] else 1)

if __name__ == "__main__":
    main()
'''


class UnsafePidExecutor:
    """Executor for Python scripts without nsjail, using subprocess isolation."""

    def __init__(
        self,
        cache_dir: str = TRACECAT__SANDBOX_CACHE_DIR,
    ):
        self.cache_dir = Path(cache_dir)
        self.package_cache = self.cache_dir / "unsafe-pid-packages"
        self.uv_cache = self.cache_dir / "uv-cache"
        self._pid_namespace_available: bool | None = None
        self._pid_namespace_probe_error: str | None = None
        self._pid_isolation_warning_emitted = False
        self._network_isolation_warning_emitted = False

        self.package_cache.mkdir(parents=True, exist_ok=True)
        self.uv_cache.mkdir(parents=True, exist_ok=True)

    def _compute_cache_key(
        self,
        dependencies: list[str],
        workspace_id: str | None = None,
    ) -> str:
        normalized = sorted(dep.lower().strip() for dep in dependencies)
        if workspace_id:
            hash_input = f"{workspace_id}\n" + "\n".join(normalized)
        else:
            hash_input = "\n".join(normalized)
        return hashlib.sha256(hash_input.encode()).hexdigest()[:16]

    @staticmethod
    def _venv_site_packages(venv_path: Path) -> Path | None:
        for candidate in sorted((venv_path / "lib").glob("python*/site-packages")):
            if candidate.exists():
                return candidate
        return None

    def _with_venv_site_packages_pythonpath(
        self,
        env_vars: dict[str, str] | None,
        venv_path: Path,
    ) -> dict[str, str]:
        merged = dict(env_vars or {})
        site_packages = self._venv_site_packages(venv_path)
        if site_packages is None:
            return merged

        pythonpath_parts = [str(site_packages)]
        if existing_pythonpath := merged.get("PYTHONPATH"):
            pythonpath_parts.append(existing_pythonpath)
        merged["PYTHONPATH"] = os.pathsep.join(pythonpath_parts)
        return merged

    def _with_python_paths(
        self,
        env_vars: dict[str, str] | None,
        python_path_dirs: list[Path],
    ) -> dict[str, str]:
        merged = dict(env_vars or {})
        pythonpath_parts = [str(path) for path in python_path_dirs if path.exists()]
        if not pythonpath_parts:
            return merged
        if existing_pythonpath := merged.get("PYTHONPATH"):
            pythonpath_parts.append(existing_pythonpath)
        merged["PYTHONPATH"] = os.pathsep.join(pythonpath_parts)
        return merged

    async def _is_pid_namespace_available(self) -> bool:
        if self._pid_namespace_available is not None:
            return self._pid_namespace_available

        if shutil.which("unshare") is None:
            self._pid_namespace_probe_error = "unshare binary not found"
            self._pid_namespace_available = False
            return False

        probe: asyncio.subprocess.Process | None = None
        try:
            probe = await asyncio.create_subprocess_exec(
                "unshare",
                "--pid",
                "--fork",
                "--kill-child",
                "true",
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
            await asyncio.wait_for(probe.wait(), timeout=2)
            self._pid_namespace_available = probe.returncode == 0
            if not self._pid_namespace_available:
                self._pid_namespace_probe_error = (
                    f"unshare probe exited with status {probe.returncode}"
                )
        except TimeoutError:
            if probe is not None:
                with suppress(ProcessLookupError):
                    probe.kill()
                await probe.wait()
            self._pid_namespace_probe_error = "unshare probe timed out"
            self._pid_namespace_available = False
        except Exception as e:
            self._pid_namespace_probe_error = f"unshare probe failed: {e}"
            self._pid_namespace_available = False
        return self._pid_namespace_available

    async def _build_execution_cmd(
        self, python_path: str, wrapper_path: Path
    ) -> list[str]:
        base_cmd = [python_path, str(wrapper_path)]
        if await self._is_pid_namespace_available():
            return ["unshare", "--pid", "--fork", "--kill-child", *base_cmd]

        if not self._pid_isolation_warning_emitted:
            message = "PID namespace isolation unavailable; running script without PID isolation"
            logger.warning(message, reason=self._pid_namespace_probe_error)
            module_logger.warning(message)
            self._pid_isolation_warning_emitted = True
        return base_cmd

    async def _create_venv(self, venv_path: Path) -> None:
        create_cmd = ["uv", "venv", str(venv_path), "--python", "3.12"]
        process = await asyncio.create_subprocess_exec(
            *create_cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env={
                "PATH": os.environ.get("PATH", ""),
                "HOME": os.environ.get("HOME", "/tmp"),
                "UV_CACHE_DIR": str(self.uv_cache),
            },
        )
        try:
            _, stderr = await asyncio.wait_for(process.communicate(), timeout=60)
        except TimeoutError as e:
            process.kill()
            await process.wait()
            raise PackageInstallError("Virtual environment creation timed out") from e
        if process.returncode != 0:
            raise PackageInstallError(
                f"Failed to create virtual environment: {stderr.decode()}"
            )

    async def _install_packages(
        self,
        venv_path: Path,
        dependencies: list[str],
        timeout_seconds: int = 300,
    ) -> None:
        pip_cmd = [
            "uv",
            "pip",
            "install",
            "--python",
            str(venv_path / "bin" / "python"),
        ]

        if TRACECAT__SANDBOX_PYPI_INDEX_URL:
            pip_cmd.extend(["--index-url", TRACECAT__SANDBOX_PYPI_INDEX_URL])
        for extra_url in TRACECAT__SANDBOX_PYPI_EXTRA_INDEX_URLS:
            pip_cmd.extend(["--extra-index-url", extra_url])
        pip_cmd.extend(dependencies)

        process = await asyncio.create_subprocess_exec(
            *pip_cmd,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env={
                "PATH": os.environ.get("PATH", ""),
                "HOME": os.environ.get("HOME", "/tmp"),
                "UV_CACHE_DIR": str(self.uv_cache),
            },
        )

        try:
            _, stderr = await asyncio.wait_for(
                process.communicate(),
                timeout=timeout_seconds,
            )
        except TimeoutError as e:
            process.kill()
            await process.wait()
            raise PackageInstallError(
                f"Package installation timed out after {timeout_seconds}s"
            ) from e

        if process.returncode != 0:
            raise PackageInstallError(f"Failed to install packages: {stderr.decode()}")

    async def execute(
        self,
        script: str,
        inputs: dict[str, Any] | None = None,
        dependencies: list[str] | None = None,
        timeout_seconds: int | None = None,
        allow_network: bool = False,
        env_vars: dict[str, str] | None = None,
        python_path_dirs: list[Path] | None = None,
        workspace_id: str | None = None,
    ) -> SandboxResult:
        if timeout_seconds is None:
            timeout_seconds = TRACECAT__SANDBOX_DEFAULT_TIMEOUT

        start_time = time.time()
        if not allow_network and not self._network_isolation_warning_emitted:
            message = "Network isolation is not enforced without nsjail; scripts may still access network"
            logger.warning(message)
            module_logger.warning(message)
            self._network_isolation_warning_emitted = True

        work_dir = Path(tempfile.mkdtemp(prefix="unsafe-pid-sandbox-"))

        try:
            python_path = shutil.which("python3") or "python3"
            execution_env_vars = self._with_python_paths(
                env_vars,
                python_path_dirs or [],
            )
            if dependencies:
                cache_key = self._compute_cache_key(dependencies, workspace_id)
                cached_venv = self.package_cache / cache_key
                if not (
                    cached_venv.exists() and (cached_venv / "bin" / "python").exists()
                ):
                    temp_venv = self.package_cache / f"{cache_key}.{os.getpid()}.tmp"
                    try:
                        if temp_venv.exists():
                            shutil.rmtree(temp_venv, ignore_errors=True)
                        await self._create_venv(temp_venv)
                        await self._install_packages(
                            temp_venv,
                            dependencies,
                            timeout_seconds=timeout_seconds,
                        )
                        try:
                            os.rename(temp_venv, cached_venv)
                        except OSError:
                            logger.debug(
                                "Venv cache race: using existing venv",
                                cache_key=cache_key,
                            )
                    finally:
                        if temp_venv.exists():
                            shutil.rmtree(temp_venv, ignore_errors=True)

                python_path = str(cached_venv / "bin" / "python")
                execution_env_vars = self._with_venv_site_packages_pythonpath(
                    execution_env_vars,
                    cached_venv,
                )

            (work_dir / "script.py").write_text(script)
            (work_dir / "inputs.json").write_text(json.dumps(inputs or {}))
            wrapper_path = work_dir / "wrapper.py"
            wrapper_path.write_text(SAFE_WRAPPER_SCRIPT.format(work_dir=str(work_dir)))

            exec_env = {
                "PATH": os.environ.get("PATH", ""),
                "HOME": os.environ.get("HOME", "/tmp"),
                "PYTHONDONTWRITEBYTECODE": "1",
                "PYTHONUNBUFFERED": "1",
                "LANG": "C.UTF-8",
                "LC_ALL": "C.UTF-8",
            }
            if execution_env_vars:
                exec_env.update(execution_env_vars)

            cmd = await self._build_execution_cmd(python_path, wrapper_path)
            process = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                cwd=str(work_dir),
                env=exec_env,
            )
            try:
                stdout_bytes, stderr_bytes = await asyncio.wait_for(
                    process.communicate(),
                    timeout=timeout_seconds,
                )
            except TimeoutError as e:
                process.kill()
                await process.wait()
                raise SandboxTimeoutError(
                    f"Script execution timed out after {timeout_seconds}s"
                ) from e

            execution_time_ms = (time.time() - start_time) * 1000
            stdout = stdout_bytes.decode("utf-8", errors="replace")
            stderr = stderr_bytes.decode("utf-8", errors="replace")

            result_path = work_dir / "result.json"
            if result_path.exists():
                try:
                    result_data = json.loads(result_path.read_text())
                    return SandboxResult(
                        success=result_data.get("success", False),
                        output=result_data.get("output"),
                        stdout=result_data.get("stdout", stdout),
                        stderr=result_data.get("stderr", stderr),
                        error=result_data.get("error"),
                        exit_code=process.returncode,
                        execution_time_ms=execution_time_ms,
                    )
                except json.JSONDecodeError:
                    logger.warning("Failed to parse result.json")

            return SandboxResult(
                success=False,
                error=f"Execution failed: {stderr[:500] if stderr else 'Unknown error'}",
                stdout=stdout,
                stderr=stderr[:500] if stderr else "",
                exit_code=process.returncode,
                execution_time_ms=execution_time_ms,
            )

        except (SandboxTimeoutError, PackageInstallError):
            raise
        except Exception as e:
            logger.error(
                "Unexpected error in unsafe PID executor",
                error_type=type(e).__name__,
            )
            raise SandboxExecutionError(
                f"Unexpected error: {type(e).__name__}: {e}"
            ) from e
        finally:
            shutil.rmtree(work_dir, ignore_errors=True)
