"""
BIBIVIDI — Script Runner
Executes BSS skill.py scripts in a sandboxed subprocess with timeout.
The script receives the prompt via stdin and must print its result to stdout.
"""
from __future__ import annotations

import asyncio
import logging
import sys
import textwrap
from dataclasses import dataclass

from bvd.config import settings
from bvd.engine.library import Script
from bvd.engine.sandbox import SECCOMP_SETUP_CODE, get_preexec_fn

logger = logging.getLogger(__name__)

# Maximum bytes read from a script's stdout (prevents OOM from runaway scripts)
MAX_STDOUT_BYTES = 512 * 1024  # 512 KB

# Modules explicitly ALLOWED inside sandboxed scripts
ALLOWED_MODULES = frozenset({
    "json", "re", "math", "statistics", "datetime", "collections",
    "itertools", "functools", "pathlib", "string", "textwrap",
    "csv", "io", "base64", "hashlib", "uuid",
    "heapq",   # needed by collections.Counter.most_common() (lazy import)
})

SANDBOX_WRAPPER = """
# ── Layer 2: seccomp BPF (Linux x86_64 only, no-op elsewhere) ──────────────
{seccomp_setup}
# ── Layer 1b: Python import allowlist ──────────────────────────────────────
import sys
import builtins

# Block dangerous imports — only top-level modules in ALLOWED are accessible.
# Transitive imports of an allowed module (e.g. re → sre_compile) are permitted
# by temporarily restoring the original import during the top-level import.
_orig_import = builtins.__import__

ALLOWED = {allowed}

def _safe_import(name, *args, **kwargs):
    top = name.split('.')[0]
    if top not in ALLOWED:
        raise ImportError(f"BIBIVIDI sandbox: import '{{name}}' is not allowed")
    # Restore original import for the duration of this call so that transitive
    # dependencies of whitelisted modules (e.g. sre_compile for re) can load.
    builtins.__import__ = _orig_import
    try:
        return _orig_import(name, *args, **kwargs)
    finally:
        builtins.__import__ = _safe_import

builtins.__import__ = _safe_import

# Inject prompt from stdin
PROMPT = sys.stdin.read()

# ── User script below ──────────────────────────────────────
{code}
"""


@dataclass
class RunResult:
    output: str
    stdout: str
    stderr: str
    exit_code: int
    timed_out: bool = False


async def run_script(script: Script, prompt: str) -> RunResult:
    """
    Execute a BSS script in a sandboxed subprocess.
    The script reads the prompt from PROMPT (injected by the wrapper).
    Its stdout becomes the condensed context sent to the LLM.

    Pending skills (status='pending') are never executed — the original
    prompt is returned so the request passes through unchanged.
    Use 'bibividi approve <id>' to activate a pending skill.
    """
    if script.status == "pending":
        logger.info(
            "Skill %s is pending approval — passing through. "
            "Run 'bibividi approve %s' to activate.",
            script.id, script.id,
        )
        return RunResult(output=prompt, stdout="", stderr="PENDING", exit_code=0)

    wrapped = _wrap(script.code)
    preexec_fn = get_preexec_fn()

    try:
        proc = await asyncio.create_subprocess_exec(
            sys.executable,
            "-c",
            wrapped,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            **({"preexec_fn": preexec_fn} if preexec_fn is not None else {}),  # type: ignore[arg-type]
        )

        try:
            stdout_b, stderr_b = await asyncio.wait_for(
                proc.communicate(input=prompt.encode()),
                timeout=settings.max_script_runtime_s,
            )
        except TimeoutError:
            proc.kill()
            await proc.communicate()
            logger.warning(
                "Script %s timed out after %.1fs", script.id, settings.max_script_runtime_s
            )
            return RunResult(
                output=prompt,   # fall back to original prompt
                stdout="",
                stderr="TIMEOUT",
                exit_code=-1,
                timed_out=True,
            )

        if len(stdout_b) > MAX_STDOUT_BYTES:
            logger.warning(
                "Script %s stdout truncated: %d bytes → %d bytes cap",
                script.id, len(stdout_b), MAX_STDOUT_BYTES,
            )
            stdout_b = stdout_b[:MAX_STDOUT_BYTES]

        stdout = stdout_b.decode(errors="replace").strip()
        stderr = stderr_b.decode(errors="replace").strip()
        exit_code = proc.returncode or 0

        if exit_code != 0:
            logger.warning(
                "Script %s exited %d: %s", script.id, exit_code, stderr[:200]
            )
            return RunResult(output=prompt, stdout=stdout, stderr=stderr, exit_code=exit_code)

        output = stdout if stdout else prompt
        logger.debug("Script %s produced %d chars", script.id, len(output))
        return RunResult(output=output, stdout=stdout, stderr=stderr, exit_code=0)

    except Exception as exc:
        logger.error("Unexpected runner error for %s: %s", script.id, exc)
        return RunResult(output=prompt, stdout="", stderr=str(exc), exit_code=-1)


def _wrap(code: str) -> str:
    """Inject the sandbox wrapper around user code."""
    return SANDBOX_WRAPPER.format(
        seccomp_setup=SECCOMP_SETUP_CODE,
        allowed=repr(ALLOWED_MODULES),
        code=textwrap.indent(code, ""),
    )


# ── BSS 1.1 Chain runner ──────────────────────────────────────────────────────

@dataclass
class ChainRunResult:
    output: str
    step_outputs: list[str]
    exit_code: int
    timed_out: bool = False
    failed_step: int | None = None  # index of the step that failed, if any


async def run_chain(
    chain_script: Script,
    prompt: str,
    step_scripts: list[Script | None],
    *,
    bypass_feature_flag: bool = False,
) -> ChainRunResult:
    """
    Execute a BSS 1.1 chain sequentially.

    Each step's stdout becomes the next step's stdin (raw string, no envelope).
    Empty stdout: the previous input passes unchanged to the next step.
    If any step fails or the total timeout fires, falls back to the original prompt.

    Requires BIBIVIDI_ENABLE_CHAINS=true; otherwise returns the original prompt.
    Set ``bypass_feature_flag=True`` for BSS 2.0 auto-assembled pipelines, which
    are governed by BIBIVIDI_ENABLE_AUTO_ASSEMBLY instead.
    """
    from bvd.config import settings

    if not settings.enable_chains and not bypass_feature_flag:
        logger.info("Chains disabled — pass-through for chain %s", chain_script.id)
        return ChainRunResult(output=prompt, step_outputs=[], exit_code=0)

    # ── Validate deps ─────────────────────────────────────────
    for i, step in enumerate(step_scripts):
        if step is None:
            logger.warning(
                "Chain %s: step %d dep not found — aborting", chain_script.id, i
            )
            return ChainRunResult(
                output=prompt, step_outputs=[], exit_code=0, failed_step=i
            )

    for i, step in enumerate(step_scripts):
        assert step is not None
        if step.status == "pending":
            logger.warning(
                "Chain %s: step %d '%s' is pending — aborting",
                chain_script.id, i, step.id,
            )
            return ChainRunResult(
                output=prompt, step_outputs=[], exit_code=0, failed_step=i
            )

    # ── Runtime cycle detection ───────────────────────────────
    visited: set[str] = {chain_script.id}
    for i, step in enumerate(step_scripts):
        assert step is not None
        if step.id in visited:
            logger.warning(
                "Chain %s: cycle detected at step %d ('%s') — aborting",
                chain_script.id, i, step.id,
            )
            return ChainRunResult(output=prompt, step_outputs=[], exit_code=0)
        visited.add(step.id)

    # ── Execute with total timeout ────────────────────────────
    try:
        return await asyncio.wait_for(
            _execute_chain(chain_script, prompt, step_scripts),  # type: ignore[arg-type]
            timeout=settings.max_chain_runtime_s,
        )
    except TimeoutError:
        logger.warning(
            "Chain %s timed out after %.1fs", chain_script.id, settings.max_chain_runtime_s
        )
        return ChainRunResult(output=prompt, step_outputs=[], exit_code=-1, timed_out=True)


async def _execute_chain(
    chain_script: Script,
    prompt: str,
    step_scripts: list[Script],
) -> ChainRunResult:
    """Inner sequential execution. Assumes validation is already done."""
    step_outputs: list[str] = []
    current_input = prompt

    # Build ordered list: declared steps + optional finalizer (chain_script.code)
    all_steps: list[Script] = list(step_scripts)
    if chain_script.code.strip():
        all_steps.append(chain_script)

    for i, step in enumerate(all_steps):
        result = await run_script(step, current_input)

        if result.timed_out:
            return ChainRunResult(
                output=prompt,
                step_outputs=step_outputs,
                exit_code=-1,
                timed_out=True,
                failed_step=i,
            )

        if result.exit_code != 0:
            logger.warning(
                "Chain %s step %d ('%s') exited %d — aborting",
                chain_script.id, i, step.id, result.exit_code,
            )
            return ChainRunResult(
                output=prompt,
                step_outputs=step_outputs,
                exit_code=result.exit_code,
                failed_step=i,
            )

        # Empty stdout → pass previous input unchanged to next step
        next_input = result.stdout.strip() if result.stdout else current_input
        step_outputs.append(next_input)
        current_input = next_input

    final_output = current_input

    # Safety: fall back if chain output is larger than the original prompt
    if len(final_output) > len(prompt):
        logger.warning(
            "Chain %s output (%d chars) > original (%d chars) — falling back",
            chain_script.id, len(final_output), len(prompt),
        )
        return ChainRunResult(output=prompt, step_outputs=step_outputs, exit_code=0)

    return ChainRunResult(output=final_output, step_outputs=step_outputs, exit_code=0)
