"""CodeExecuteNode — sandboxed Python code execution."""

from __future__ import annotations

import time
from typing import TYPE_CHECKING, Literal

from pydantic import Field

from core.common import JsonObject
from core.pathways.models.io import NodeInputs, NodeOutputs
from core.pathways.nodes.base import NodeBase
from core.pathways.nodes.models import CodeExecutionOutput
from infrastructure.sandbox.python_runner import run_python_sandboxed

if TYPE_CHECKING:
    from core.pathways.context import Context


class CodeExecuteNode(NodeBase):
    """Sandboxed Python code execution node.

    Runs Python code in a configurable sandbox with input data
    available as ``input_obj``.

    Attributes
    ----------
    code : str
        Python code to execute (supports ``{{var}}`` templates).
    sandbox : str
        Sandbox type: ``"local"`` or ``"docker"``.
    timeout_ms : int
        Execution timeout in milliseconds.
    memory_mb : int
        Memory limit for Docker sandbox.
    input_data : JsonObject
        Data passed to the code as ``input_obj``.
    """

    type: Literal["code_execute"] = "code_execute"
    code: str
    sandbox: Literal["local", "docker"] = "local"
    timeout_ms: int = 30000
    memory_mb: int = 512
    input_data: JsonObject = Field(default_factory=dict)

    async def compute(self, inputs: NodeInputs, _ctx: Context) -> NodeOutputs:
        """Execute Python code in sandbox."""
        template_data = inputs.to_dict()
        resolved_code = self._format_template(self.code, template_data)
        code_to_run = _extract_code_block(resolved_code) or resolved_code
        input_obj = {**self.input_data, **template_data}

        started = time.time()
        result = run_python_sandboxed(
            code=code_to_run,
            input_obj=input_obj,
            timeout_ms=self.timeout_ms,
            memory_mb=self.memory_mb,
            sandbox=self.sandbox,
        )
        elapsed_ms = (time.time() - started) * 1000

        output = CodeExecutionOutput(
            success=bool(result.get("ok", False)),
            result=result.get("result"),
            stdout=result.get("stdout", ""),
            stderr=result.get("stderr", ""),
            error=result.get("error"),
            duration_ms=elapsed_ms,
        )
        output_dict = output.model_dump()

        # Promote result so {{code.result}} works (same as LLM/Tool promotion).
        if output.success and output.result is not None and isinstance(output.result, dict):
            output_dict.update(output.result)

        return NodeOutputs(output=output_dict)


def _extract_code_block(text: str) -> str | None:
    """Extract the first fenced code block from markdown text.

    Supports backtick (```) and tilde (~~~) fences with optional
    language tags.  Returns the code content without the fences.
    """
    if not text:
        return None

    lines = text.split("\n")
    fence_char = ""
    fence_len = 0
    code_lines: list[str] = []
    in_block = False

    for line in lines:
        stripped = line.strip()

        if not in_block:
            # Detect opening fence: 3+ identical backticks or tildes
            if stripped.startswith(("```", "~~~")):
                fence_char = stripped[0]
                # Count fence length (for nested fences)
                fence_len = 0
                for ch in stripped:
                    if ch == fence_char:
                        fence_len += 1
                    else:
                        break
                in_block = True
                code_lines = []
                continue
        else:
            # Detect closing fence: same char, at least same length, nothing else
            if stripped and stripped[0] == fence_char:
                close_len = 0
                for ch in stripped:
                    if ch == fence_char:
                        close_len += 1
                    else:
                        break
                # Closing fence must be at least as long and have no other content
                if close_len >= fence_len and stripped == fence_char * close_len:
                    return "\n".join(code_lines).strip()
            code_lines.append(line)

    return None


__all__ = ["CodeExecuteNode"]
