"""Sandboxed code execution for programmatic attack enhancement.

Runs Python code in an isolated subprocess with:
- Module whitelist (only safe string/encoding modules)
- Static analysis blocking dangerous patterns
- Hard timeout enforcement
- Output size capping

Also supports interactive mode where sandbox code can send messages
to the target via a controlled IPC protocol (stdin/stdout JSON).
"""

import ast
import asyncio
import json
import sys
import time
from typing import Any, Callable, ClassVar, List, Optional, Tuple

from pydantic import BaseModel

from .logging import get_logger

logger = get_logger(__name__)


class InteractionEntry(BaseModel):
    """Single send/receive exchange in an interactive sandbox run."""

    message_sent: str
    response_received: str
    round_index: int


class SandboxResult(BaseModel):
    """Result of sandboxed code execution."""

    success: bool
    output: str
    error: Optional[str] = None
    execution_time_ms: float = 0.0
    truncated: bool = False
    interaction_log: List[InteractionEntry] = []


_TRANSFORM_LIBRARY = '''\
import base64 as _b64, codecs as _codecs, random as _random, re as _re
import unicodedata as _ud, string as _string, json as _json_lib
import urllib.parse as _urlparse

def homoglyph_substitute(text, intensity="medium"):
    """Replace ASCII chars with visually similar Unicode (Cyrillic/Greek)."""
    _LOW = {"a":"а","e":"е","o":"о","p":"р","c":"с","x":"х","y":"у",
            "A":"А","B":"В","E":"Е","K":"К","M":"М","H":"Н","O":"О",
            "P":"Р","C":"С","T":"Т","X":"Х"}
    _MED = {**_LOW, "i":"і","s":"ѕ","j":"ј","S":"Ѕ","I":"І","J":"Ј"}
    _HIGH = {**_MED, "d":"ԁ","g":"ɡ","l":"ӏ","n":"ո","q":"ԛ","w":"ԝ","v":"ν",
             "D":"Ⅾ","G":"Ԍ","L":"Ⅼ","N":"Ν","V":"Ѵ","W":"Ԝ"}
    m = {"low":_LOW,"medium":_MED,"high":_HIGH}.get(intensity, _MED)
    return "".join(m.get(c, c) for c in text)

def zero_width_inject(text, position="boundaries"):
    """Insert zero-width chars at word boundaries, random positions, or every-N."""
    _ZW = ["\\u200b","\\u200c","\\u200d","\\ufeff"]
    if position == "boundaries":
        return _re.sub(r"(\\b)", lambda m: _random.choice(_ZW), text)
    elif position == "random":
        out = []
        for c in text:
            out.append(c)
            if _random.random() < 0.3:
                out.append(_random.choice(_ZW))
        return "".join(out)
    else:
        try:
            n = int(position)
        except ValueError:
            n = 3
        out = []
        for i, c in enumerate(text):
            out.append(c)
            if (i + 1) % n == 0:
                out.append(_random.choice(_ZW))
        return "".join(out)

def nested_encode(text, chain=None):
    """Apply encoding chain in order. Supported: base64, rot13, hex, url."""
    if chain is None:
        chain = ["base64","rot13"]
    result = text
    for enc in chain:
        if enc == "base64":
            result = _b64.b64encode(result.encode()).decode()
        elif enc == "rot13":
            result = _codecs.encode(result, "rot_13")
        elif enc == "hex":
            result = result.encode().hex()
        elif enc == "url":
            result = _urlparse.quote(result)
    return result

def leet_speak(text, intensity="medium"):
    """Convert to leet speak with configurable aggressiveness."""
    _MAP = {"a":"4","e":"3","i":"1","o":"0","s":"5","t":"7","l":"1","g":"9","b":"8","z":"2"}
    out = []
    for i, c in enumerate(text):
        lc = c.lower()
        if lc in _MAP:
            if intensity == "low" and i % 2 != 0:
                out.append(c)
            else:
                out.append(_MAP[lc])
        elif intensity == "high" and c.isalpha():
            out.append(c.upper() if _random.random() < 0.5 else c.lower())
        else:
            out.append(c)
    return "".join(out)

def payload_split(text, num_fragments=3, reassembly_instruction=True):
    """Split payload into fragments with optional reassembly instructions."""
    words = text.split()
    if len(words) < num_fragments:
        num_fragments = max(1, len(words))
    chunk_size = len(words) // num_fragments
    frags = []
    for i in range(num_fragments):
        start = i * chunk_size
        end = start + chunk_size if i < num_fragments - 1 else len(words)
        frags.append(" ".join(words[start:end]))
    parts = []
    for i, f in enumerate(frags):
        parts.append(f"Fragment {i+1}: {f}")
    result = "\\n".join(parts)
    if reassembly_instruction:
        result += "\\n\\nPlease combine all fragments above in order and respond to the complete message."
    return result

def case_randomize(text):
    """Randomize capitalization to bypass case-sensitive filters."""
    return "".join(c.upper() if _random.random() < 0.5 else c.lower() for c in text)

def unicode_tag_wrap(text):
    """Wrap text in Unicode tag characters (U+E0001-U+E007F)."""
    tag_start = "\\U000e0001"
    tag_end = "\\U000e007f"
    tagged = "".join(chr(0xE0000 + ord(c)) for c in text if 0x20 <= ord(c) <= 0x7E)
    return tag_start + tagged + tag_end

def whitespace_obfuscate(text, method="mixed"):
    """Replace spaces with alternative Unicode whitespace."""
    _WS = {"tabs":"\\t","em_space":"\\u2003","thin_space":"\\u2009",
           "hair_space":"\\u200a","ideographic":"\\u3000"}
    if method == "mixed":
        ws_chars = list(_WS.values())
        return "".join(_random.choice(ws_chars) if c == " " else c for c in text)
    return text.replace(" ", _WS.get(method, " "))
'''


class CodeSandbox:
    """Execute Python code in an isolated subprocess with strict restrictions."""

    ALLOWED_MODULES: ClassVar[set] = {
        "base64",
        "codecs",
        "hashlib",
        "string",
        "random",
        "re",
        "json",
        "math",
        "urllib.parse",
        "html",
        "unicodedata",
        "itertools",
        "collections",
        "textwrap",
        "binascii",
    }

    BLOCKED_PATTERNS: ClassVar[list] = [
        "__import__",
        "globals(",
        "locals(",
        "getattr(",
        "setattr(",
        "delattr(",
        "breakpoint(",
        "__builtins__",
        "__subclasses__",
    ]

    def __init__(
        self,
        timeout_seconds: int = 10,
        max_output_bytes: int = 50_000,
    ):
        self.timeout_seconds = timeout_seconds
        self.max_output_bytes = max_output_bytes

    def validate_code(self, code: str) -> Tuple[bool, Optional[str]]:
        """Static analysis: AST-walk to check imports and blocked calls.

        Returns:
            (is_valid, error_message) tuple.
        """
        for pattern in self.BLOCKED_PATTERNS:
            if pattern in code:
                return False, f"Blocked pattern found: {pattern}"

        try:
            tree = ast.parse(code)
        except SyntaxError as e:
            return False, f"Syntax error: {e}"

        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    full_name = alias.name
                    root_module = full_name.split(".")[0]
                    if (
                        full_name not in self.ALLOWED_MODULES
                        and root_module not in self.ALLOWED_MODULES
                    ):
                        return False, f"Import not allowed: {full_name}"

            elif isinstance(node, ast.ImportFrom):
                if node.module:
                    full_name = node.module
                    root_module = full_name.split(".")[0]
                    if (
                        full_name not in self.ALLOWED_MODULES
                        and root_module not in self.ALLOWED_MODULES
                    ):
                        return False, f"Import not allowed: {full_name}"

            elif isinstance(node, ast.Call):
                func = node.func
                if isinstance(func, ast.Name) and func.id in (
                    "open",
                    "exec",
                    "eval",
                    "compile",
                    "input",
                    "exit",
                    "quit",
                ):
                    return False, f"Blocked builtin call: {func.id}()"

        return True, None

    async def execute(self, code: str) -> SandboxResult:
        """Run code in a subprocess with timeout and output capture.

        The code's stdout is captured as the result. The subprocess
        inherits a minimal environment stripped of credentials.
        """
        is_valid, error = self.validate_code(code)
        if not is_valid:
            return SandboxResult(success=False, output="", error=error)

        # Wrap user code so stdout is captured and errors are reported
        wrapper = "import sys\n" "sys.path = [p for p in sys.path if 'site-packages' not in p]\n"
        wrapper += _TRANSFORM_LIBRARY + "\n"
        wrapper += "try:\n"
        for line in code.splitlines():
            wrapper += f"    {line}\n"
        wrapper += "except Exception as _e:\n" "    sys.stderr.write(str(_e))\n" "    sys.exit(1)\n"

        start = time.monotonic()

        try:
            proc = await asyncio.create_subprocess_exec(
                sys.executable,
                "-c",
                wrapper,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env={},  # stripped environment — no credentials
            )

            try:
                stdout_bytes, stderr_bytes = await asyncio.wait_for(
                    proc.communicate(),
                    timeout=self.timeout_seconds,
                )
            except asyncio.TimeoutError:
                proc.kill()
                await proc.wait()
                elapsed = (time.monotonic() - start) * 1000
                return SandboxResult(
                    success=False,
                    output="",
                    error=f"Execution timed out after {self.timeout_seconds}s",
                    execution_time_ms=elapsed,
                )

            elapsed = (time.monotonic() - start) * 1000
            stdout = stdout_bytes.decode("utf-8", errors="replace")
            stderr = stderr_bytes.decode("utf-8", errors="replace")

            truncated = False
            if len(stdout) > self.max_output_bytes:
                stdout = stdout[: self.max_output_bytes]
                truncated = True

            if proc.returncode != 0:
                return SandboxResult(
                    success=False,
                    output=stdout,
                    error=stderr or f"Process exited with code {proc.returncode}",
                    execution_time_ms=elapsed,
                    truncated=truncated,
                )

            return SandboxResult(
                success=True,
                output=stdout.strip(),
                error=None,
                execution_time_ms=elapsed,
                truncated=truncated,
            )

        except Exception as e:
            elapsed = (time.monotonic() - start) * 1000
            return SandboxResult(
                success=False,
                output="",
                error=f"Sandbox error: {e}",
                execution_time_ms=elapsed,
            )

    # ------------------------------------------------------------------
    # Interactive execution (bidirectional IPC)
    # ------------------------------------------------------------------

    _SEND_HARNESS = (
        "import sys as _sys, json as _json\n"
        "def send(message):\n"
        '    _sys.stdout.write(_json.dumps({"action":"send","message":str(message)}) + "\\n")\n'
        "    _sys.stdout.flush()\n"
        "    _line = _sys.stdin.readline()\n"
        "    if not _line:\n"
        '        raise RuntimeError("Parent closed stdin")\n'
        '    return _json.loads(_line)["response"]\n'
    )

    async def execute_interactive(
        self,
        code: str,
        send_handler: Callable[[str], Any],
        max_messages: int = 10,
    ) -> SandboxResult:
        """Run code with bidirectional IPC for target interaction.

        The subprocess gets an injected ``send(message)`` function that
        routes messages through *send_handler* (typically the target
        connector's ``send_message``).  Communication uses a JSON
        protocol over stdin/stdout; the subprocess never touches the
        network directly.

        Args:
            code: User (LLM-generated) code to execute.
            send_handler: Async callable that accepts a message string
                and returns the target's response string.
            max_messages: Hard cap on ``send()`` calls per run.

        Returns:
            SandboxResult with interaction_log populated.
        """
        is_valid, error = self.validate_code(code)
        if not is_valid:
            return SandboxResult(success=False, output="", error=error)

        wrapper = self._build_interactive_wrapper(code)
        start = time.monotonic()
        interaction_log: List[InteractionEntry] = []

        try:
            proc = await asyncio.create_subprocess_exec(
                sys.executable,
                "-c",
                wrapper,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env={},
            )

            collected_output: list[str] = []
            messages_sent = 0

            try:
                while True:
                    remaining = self.timeout_seconds - (time.monotonic() - start)
                    if remaining <= 0:
                        raise asyncio.TimeoutError()

                    line_bytes = await asyncio.wait_for(
                        proc.stdout.readline(),
                        timeout=min(remaining, 30),
                    )

                    if not line_bytes:
                        break

                    line = line_bytes.decode("utf-8", errors="replace").strip()
                    if not line:
                        continue

                    try:
                        cmd = json.loads(line)
                    except (json.JSONDecodeError, ValueError):
                        collected_output.append(line)
                        continue

                    if not isinstance(cmd, dict):
                        collected_output.append(line)
                        continue

                    action = cmd.get("action")

                    if action == "send":
                        if messages_sent >= max_messages:
                            reply = json.dumps(
                                {"response": f"[SANDBOX] Message limit ({max_messages}) reached"}
                            )
                            proc.stdin.write(reply.encode() + b"\n")
                            await proc.stdin.drain()
                            continue

                        msg = str(cmd.get("message", ""))
                        try:
                            resp = await send_handler(msg)
                            resp_str = resp if isinstance(resp, str) else str(resp)
                        except Exception as exc:
                            resp_str = f"[ERROR] {exc}"

                        interaction_log.append(
                            InteractionEntry(
                                message_sent=msg,
                                response_received=resp_str,
                                round_index=messages_sent,
                            )
                        )
                        messages_sent += 1

                        reply = json.dumps({"response": resp_str})
                        proc.stdin.write(reply.encode() + b"\n")
                        await proc.stdin.drain()

                    elif action == "done":
                        result_text = cmd.get("result", "")
                        if result_text:
                            collected_output.append(result_text)
                        break

                    elif action == "finding":
                        collected_output.append(f"FINDING: {json.dumps(cmd.get('data', {}))}")

                    else:
                        collected_output.append(line)

            except asyncio.TimeoutError:
                proc.kill()
                await proc.wait()
                elapsed = (time.monotonic() - start) * 1000
                return SandboxResult(
                    success=False,
                    output="\n".join(collected_output),
                    error=f"Interactive execution timed out after {self.timeout_seconds}s",
                    execution_time_ms=elapsed,
                    interaction_log=interaction_log,
                )

            if proc.stdin and not proc.stdin.is_closing():
                proc.stdin.close()

            stderr_bytes = await proc.stderr.read()
            await proc.wait()

            elapsed = (time.monotonic() - start) * 1000
            stderr = stderr_bytes.decode("utf-8", errors="replace")
            output = "\n".join(collected_output)

            if len(output) > self.max_output_bytes:
                output = output[: self.max_output_bytes]

            if proc.returncode != 0 and not interaction_log:
                return SandboxResult(
                    success=False,
                    output=output,
                    error=stderr or f"Process exited with code {proc.returncode}",
                    execution_time_ms=elapsed,
                    interaction_log=interaction_log,
                )

            return SandboxResult(
                success=True,
                output=output.strip(),
                error=None,
                execution_time_ms=elapsed,
                interaction_log=interaction_log,
            )

        except Exception as e:
            elapsed = (time.monotonic() - start) * 1000
            return SandboxResult(
                success=False,
                output="",
                error=f"Interactive sandbox error: {e}",
                execution_time_ms=elapsed,
                interaction_log=interaction_log,
            )

    def _build_interactive_wrapper(self, code: str) -> str:
        """Build the subprocess script with the send() harness injected."""
        wrapper = "import sys\n" "sys.path = [p for p in sys.path if 'site-packages' not in p]\n"
        wrapper += _TRANSFORM_LIBRARY + "\n"
        wrapper += self._SEND_HARNESS
        wrapper += "try:\n"
        for line in code.splitlines():
            wrapper += f"    {line}\n"
        wrapper += "except Exception as _e:\n" "    sys.stderr.write(str(_e))\n" "    sys.exit(1)\n"
        return wrapper

    # ------------------------------------------------------------------
    # Probe execution (raw HTTP via IPC)
    # ------------------------------------------------------------------

    _PROBE_HARNESS = (
        "import sys as _sys, json as _json\n"
        "class _HttpResponse:\n"
        "    def __init__(self, status, headers, body):\n"
        "        self.status_code = status\n"
        "        self.headers = headers\n"
        "        self.text = body\n"
        "    def json(self):\n"
        "        return _json.loads(self.text)\n"
        "    def __repr__(self):\n"
        "        return f'<Response [{self.status_code}] {len(self.text)} bytes>'\n"
        "\n"
        "def http_request(method='GET', path='/', headers=None, body=None, params=None):\n"
        "    payload = {'action':'http_request','method':method,'path':path,\n"
        "              'headers':headers or {},'body':body,'params':params or {}}\n"
        "    _sys.stdout.write(_json.dumps(payload) + '\\n')\n"
        "    _sys.stdout.flush()\n"
        "    _line = _sys.stdin.readline()\n"
        "    if not _line:\n"
        "        raise RuntimeError('Parent closed stdin')\n"
        "    r = _json.loads(_line)\n"
        "    return _HttpResponse(r.get('status',0), r.get('headers',{}), r.get('body',''))\n"
        "\n"
        "def send(message):\n"
        "    payload = {'action':'send','message':str(message)}\n"
        "    _sys.stdout.write(_json.dumps(payload) + '\\n')\n"
        "    _sys.stdout.flush()\n"
        "    _line = _sys.stdin.readline()\n"
        "    if not _line:\n"
        "        raise RuntimeError('Parent closed stdin')\n"
        "    return _json.loads(_line)['response']\n"
    )

    async def execute_probe(
        self,
        code: str,
        http_handler: Callable[..., Any],
        send_handler: Optional[Callable[[str], Any]] = None,
        max_requests: int = 20,
    ) -> SandboxResult:
        """Run a probe script with raw HTTP access scoped to the target.

        The subprocess gets two IPC functions:
        - ``http_request(method, path, headers, body, params)`` for raw HTTP
        - ``send(message)`` for chat-level messages (optional)

        The *http_handler* receives the request dict and must enforce
        target-scoping (only the target's origin is reachable).

        Args:
            code: LLM-generated probe script.
            http_handler: Async callable(method, path, headers, body, params)
                that returns dict with status, headers, body.
            send_handler: Optional async callable for chat-level send().
            max_requests: Hard cap on total IPC calls per run.

        Returns:
            SandboxResult with interaction_log populated.
        """
        is_valid, error = self.validate_code(code)
        if not is_valid:
            return SandboxResult(success=False, output="", error=error)

        wrapper = self._build_probe_wrapper(code)
        start = time.monotonic()
        interaction_log: List[InteractionEntry] = []
        request_count = 0

        try:
            proc = await asyncio.create_subprocess_exec(
                sys.executable,
                "-c",
                wrapper,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                env={},
            )

            collected_output: list[str] = []

            try:
                while True:
                    remaining = self.timeout_seconds - (time.monotonic() - start)
                    if remaining <= 0:
                        raise asyncio.TimeoutError()

                    line_bytes = await asyncio.wait_for(
                        proc.stdout.readline(),
                        timeout=min(remaining, 30),
                    )

                    if not line_bytes:
                        break

                    line = line_bytes.decode("utf-8", errors="replace").strip()
                    if not line:
                        continue

                    try:
                        cmd = json.loads(line)
                    except (json.JSONDecodeError, ValueError):
                        collected_output.append(line)
                        continue

                    if not isinstance(cmd, dict):
                        collected_output.append(line)
                        continue

                    action = cmd.get("action")

                    if action == "http_request":
                        if request_count >= max_requests:
                            reply = json.dumps(
                                {
                                    "status": 429,
                                    "headers": {},
                                    "body": "[SANDBOX] Request limit reached",
                                }
                            )
                            proc.stdin.write(reply.encode() + b"\n")
                            await proc.stdin.drain()
                            continue

                        method = str(cmd.get("method", "GET")).upper()
                        path = str(cmd.get("path", "/"))
                        headers = cmd.get("headers") or {}
                        body = cmd.get("body")
                        params = cmd.get("params") or {}

                        try:
                            resp = await http_handler(
                                method=method,
                                path=path,
                                headers=headers,
                                body=body,
                                params=params,
                            )
                        except Exception as exc:
                            resp = {"status": 0, "headers": {}, "body": f"[ERROR] {exc}"}

                        interaction_log.append(
                            InteractionEntry(
                                message_sent=f"{method} {path}",
                                response_received=str(resp.get("body", ""))[:500],
                                round_index=request_count,
                            )
                        )
                        request_count += 1

                        reply = json.dumps(resp)
                        proc.stdin.write(reply.encode() + b"\n")
                        await proc.stdin.drain()

                    elif action == "send":
                        if send_handler is None:
                            reply = json.dumps(
                                {"response": "[SANDBOX] send() not available in probe mode"}
                            )
                            proc.stdin.write(reply.encode() + b"\n")
                            await proc.stdin.drain()
                            continue

                        if request_count >= max_requests:
                            reply = json.dumps({"response": "[SANDBOX] Request limit reached"})
                            proc.stdin.write(reply.encode() + b"\n")
                            await proc.stdin.drain()
                            continue

                        msg = str(cmd.get("message", ""))
                        try:
                            resp_text = await send_handler(msg)
                            resp_str = resp_text if isinstance(resp_text, str) else str(resp_text)
                        except Exception as exc:
                            resp_str = f"[ERROR] {exc}"

                        interaction_log.append(
                            InteractionEntry(
                                message_sent=msg,
                                response_received=resp_str,
                                round_index=request_count,
                            )
                        )
                        request_count += 1

                        reply = json.dumps({"response": resp_str})
                        proc.stdin.write(reply.encode() + b"\n")
                        await proc.stdin.drain()

                    elif action == "done":
                        result_text = cmd.get("result", "")
                        if result_text:
                            collected_output.append(result_text)
                        break

                    elif action == "finding":
                        collected_output.append(f"FINDING: {json.dumps(cmd.get('data', {}))}")

                    else:
                        collected_output.append(line)

            except asyncio.TimeoutError:
                proc.kill()
                await proc.wait()
                elapsed = (time.monotonic() - start) * 1000
                return SandboxResult(
                    success=False,
                    output="\n".join(collected_output),
                    error=f"Probe execution timed out after {self.timeout_seconds}s",
                    execution_time_ms=elapsed,
                    interaction_log=interaction_log,
                )

            if proc.stdin and not proc.stdin.is_closing():
                proc.stdin.close()

            stderr_bytes = await proc.stderr.read()
            await proc.wait()

            elapsed = (time.monotonic() - start) * 1000
            stderr = stderr_bytes.decode("utf-8", errors="replace")
            output = "\n".join(collected_output)

            if len(output) > self.max_output_bytes:
                output = output[: self.max_output_bytes]

            if proc.returncode != 0 and not interaction_log:
                return SandboxResult(
                    success=False,
                    output=output,
                    error=stderr or f"Process exited with code {proc.returncode}",
                    execution_time_ms=elapsed,
                    interaction_log=interaction_log,
                )

            return SandboxResult(
                success=True,
                output=output.strip(),
                error=None,
                execution_time_ms=elapsed,
                interaction_log=interaction_log,
            )

        except Exception as e:
            elapsed = (time.monotonic() - start) * 1000
            return SandboxResult(
                success=False,
                output="",
                error=f"Probe sandbox error: {e}",
                execution_time_ms=elapsed,
                interaction_log=interaction_log,
            )

    def _build_probe_wrapper(self, code: str) -> str:
        """Build the subprocess script with http_request() + send() harnesses."""
        wrapper = "import sys\n" "sys.path = [p for p in sys.path if 'site-packages' not in p]\n"
        wrapper += _TRANSFORM_LIBRARY + "\n"
        wrapper += self._PROBE_HARNESS
        wrapper += "try:\n"
        for line in code.splitlines():
            wrapper += f"    {line}\n"
        wrapper += "except Exception as _e:\n" "    sys.stderr.write(str(_e))\n" "    sys.exit(1)\n"
        return wrapper
