"""Sandbox environment, dependency, and transport helpers."""

from __future__ import annotations

import json
import shlex
from pathlib import Path
from typing import Any, Optional

from evaluation.models import UserSubmission
from evaluation.transport import GrpcSandboxTransport, SandboxTransport


def write_files_chunked(sandbox: Any, files: dict[str, bytes], base_dir: str) -> None:
    items = list(files.items())
    chunk_size = 200
    for i in range(0, len(items), chunk_size):
        chunk = items[i : i + chunk_size]
        payload = [{"path": f"{base_dir}/{path}", "data": data} for path, data in chunk]
        sandbox.files.write_files(payload)


def ensure_runtime_env(sandbox: Any, deps_timeout_seconds: int) -> None:
    sandbox.commands.run(
        "bash -lc \"pip install -q uv && "
        "rm -rf /workspace/.venv && "
        "uv venv /workspace/.venv\"",
        timeout=int(deps_timeout_seconds),
    )
    r = sandbox.commands.run(
        "bash -lc \"test -x /workspace/.venv/bin/python && echo OK || echo FAIL\"",
        timeout=10,
    )
    if "OK" not in (r.stdout or ""):
        raise RuntimeError(
            "venv creation failed: /workspace/.venv/bin/python not found. "
            "Check network connectivity and pip/uv availability in sandbox."
        )


def install_transport_dependencies(sandbox: Any, deps_timeout_seconds: int) -> None:
    sandbox.commands.run(
        "bash -lc \"cd /workspace && source .venv/bin/activate && "
        "uv pip install -q grpcio\"",
        timeout=int(deps_timeout_seconds),
    )


def install_problem_dependencies(
    *,
    config: Any,
    sandbox: Any,
    deps_timeout_seconds: int,
    extract_requirement_pkg_name: Any,
) -> None:
    raw = getattr(config, "pip_dependencies", [])
    if not isinstance(raw, list):
        return
    deps: list[str] = []
    for item in raw:
        dep = str(item).strip()
        if dep:
            deps.append(dep)
    if not deps:
        return

    allowed_raw = getattr(config, "allowed_packages", []) or []
    if isinstance(allowed_raw, list) and allowed_raw:
        whitelist = {
            str(pkg).lower().replace("-", "_").replace(".", "_")
            for pkg in allowed_raw
            if str(pkg).strip()
        }
        blocked: list[str] = []
        for dep in deps:
            pkg_name = extract_requirement_pkg_name(dep)
            if not pkg_name or pkg_name not in whitelist:
                blocked.append(dep)
        if blocked:
            raise ValueError(
                "phase_config.pip_dependencies contains packages not in allowed_packages whitelist: "
                + ", ".join(blocked)
            )

    install_args = " ".join(shlex.quote(dep) for dep in deps)
    sandbox.commands.run(
        f"bash -lc \"cd /workspace && source .venv/bin/activate && "
        f"uv pip install -q {install_args}\"",
        timeout=int(deps_timeout_seconds),
    )


def to_positive_int(raw: Any) -> Optional[int]:
    try:
        value = int(raw)
    except Exception:
        return None
    return value if value > 0 else None


def resolve_sandbox_resources(config: Any) -> tuple[Optional[int], Optional[int]]:
    cpu_count = to_positive_int(
        getattr(config, "sandbox_cpu_count", None) or getattr(config, "cpu_count", None)
    )
    memory_mb = to_positive_int(
        getattr(config, "memory_limit_mb", None)
        or getattr(config, "memory_limit_mib", None)
        or getattr(config, "memory_mb", None)
    )
    return cpu_count, memory_mb


def read_file_snippet(sandbox: Any, path: str, max_bytes: int = 8192) -> str:
    if max_bytes <= 0:
        return ""
    cmd = (
        "bash -lc \"python -c "
        + shlex.quote(
            "import pathlib,sys;"
            f"p=pathlib.Path({path!r});"
            f"n={int(max_bytes)};"
            "b=(p.read_bytes()[-n:] if p.exists() else b'');"
            "sys.stdout.write(b.decode('utf-8','ignore'))"
        )
        + "\""
    )
    try:
        r = sandbox.commands.run(cmd, timeout=8)
        return (r.stdout or "")[:max_bytes]
    except Exception:
        return ""


def latest_mem_used_mib(sandbox: Any) -> Optional[float]:
    getter = getattr(sandbox, "get_metrics", None)
    if not callable(getter):
        return None
    try:
        metrics = getter()
    except Exception:
        return None
    if not metrics:
        return None
    last = metrics[-1]
    if isinstance(last, dict):
        raw = (
            last.get("mem_used_mib")
            or last.get("memUsedMiB")
        )
    else:
        raw = (
            getattr(last, "mem_used_mib", None)
            or getattr(last, "memUsedMiB", None)
        )
    try:
        return float(raw)
    except Exception:
        return None


def is_likely_mle_exit(config: Any, sandbox: Any, stderr_path: str) -> bool:
    _, limit_mb = resolve_sandbox_resources(config)
    used_mib = latest_mem_used_mib(sandbox)
    if limit_mb and used_mib is not None:
        if used_mib >= float(limit_mb) * 0.98:
            return True

    stderr_text = read_file_snippet(sandbox, stderr_path, max_bytes=12000).lower()
    if not stderr_text:
        return False

    strong_keywords = (
        "memoryerror",
        "cannot allocate memory",
        "out of memory",
        "std::bad_alloc",
    )
    has_strong_signal = any(k in stderr_text for k in strong_keywords)
    if not has_strong_signal:
        return False

    if limit_mb and used_mib is not None:
        return used_mib >= float(limit_mb) * 0.85

    return True


def build_llm_env_vars(gateway_token: str, gateway_url: str) -> dict[str, str]:
    return {
        "LLM_GATEWAY_TOKEN": gateway_token,
        "LLM_GATEWAY_URL": gateway_url,
        "LLM_API_KEY": gateway_token,
        "LLM_BASE_URL": gateway_url,
        "OPENAI_API_KEY": gateway_token,
        "OPENAI_BASE_URL": gateway_url,
        "DEEPSEEK_API_KEY": gateway_token,
        "DEEPSEEK_BASE_URL": gateway_url,
        "QWEN_API_KEY": gateway_token,
        "QWEN_BASE_URL": gateway_url,
        "GROK_API_KEY": gateway_token,
        "GROK_BASE_URL": gateway_url,
        "GEMINI_API_KEY": gateway_token,
        "GEMINI_BASE_URL": gateway_url,
        "ANTHROPIC_API_KEY": gateway_token,
        "ANTHROPIC_BASE_URL": gateway_url,
        "SILICONFLOW_API_KEY": gateway_token,
        "SILICONFLOW_BASE_URL": gateway_url,
        "MOONSHOT_API_KEY": gateway_token,
        "MOONSHOT_BASE_URL": gateway_url,
        "MINIMAX_API_KEY": gateway_token,
        "MINIMAX_BASE_URL": gateway_url,
        "OPENROUTER_API_KEY": gateway_token,
        "OPENROUTER_BASE_URL": gateway_url,
    }


def build_judge_envs(
    *,
    config: Any,
    submission: UserSubmission,
    get_client: Any,
    gateway_token: Optional[str] = None,
) -> dict[str, str]:
    envs: dict[str, str] = {
        "SUBMIT_ID": str(submission.submit_id),
        "PHASE_ID": str(submission.phase_id),
        "USER_ID": str(submission.user_id),
        "PHASE_CONFIG": json.dumps(config.model_dump(), ensure_ascii=False),
        "RUNTIME_CONFIG": json.dumps(submission.runtime_config.model_dump(), ensure_ascii=False),
        "PYTHONPATH": "/workspace/judge",
        "DUAL_SANDBOX_MODE": "judge",
        "PYTHONUNBUFFERED": "1",
    }
    extra_env = getattr(config, "judge_envs", {}) or {}
    if isinstance(extra_env, dict):
        for k, v in extra_env.items():
            if k and isinstance(v, (str, int, float, bool)):
                envs[str(k)] = str(v)
    if gateway_token:
        gateway_base_url = f"{get_client().base_url}/gateway/v1"
        envs.update(build_llm_env_vars(gateway_token, gateway_base_url))
    return envs


def build_user_envs(
    *,
    config: Any,
    submission: UserSubmission,
    get_client: Any,
    gateway_token: Optional[str],
) -> dict[str, str]:
    envs: dict[str, str] = {
        "SUBMIT_ID": str(submission.submit_id),
        "PHASE_ID": str(submission.phase_id),
        "USER_ID": str(submission.user_id),
        "PHASE_CONFIG": json.dumps(config.model_dump(), ensure_ascii=False),
        "RUNTIME_CONFIG": json.dumps(submission.runtime_config.model_dump(), ensure_ascii=False),
        "PYTHONPATH": "/workspace/user",
        "DUAL_SANDBOX_MODE": "user",
        "PYTHONUNBUFFERED": "1",
    }
    if gateway_token:
        gateway_base_url = f"{get_client().base_url}/gateway/v1"
        envs.update(build_llm_env_vars(gateway_token, gateway_base_url))
    return envs


def load_grpc_bridge_support_files(base_file: str) -> dict[str, bytes]:
    root = Path(base_file).resolve().parent
    proto_base = root / "proto"
    runtime_base = root / "runtime"
    required = [
        (proto_base / "eval_bridge_pb2.py", "eval_bridge_pb2.py"),
        (proto_base / "eval_bridge_pb2_grpc.py", "eval_bridge_pb2_grpc.py"),
        (runtime_base / "judge_runtime.py", "eval_runtime/judge_runtime.py"),
        (runtime_base / "judge_scaffold.py", "eval_runtime/judge_scaffold.py"),
        (runtime_base / "user_runtime.py", "eval_runtime/user_runtime.py"),
        (runtime_base / "user_adapter.py", "eval_runtime/user_adapter.py"),
    ]
    files: dict[str, bytes] = {"eval_runtime/__init__.py": b""}
    for path, target in required:
        if not path.exists():
            raise FileNotFoundError(f"missing grpc bridge support file: {path}")
        files[target] = path.read_bytes()
    return files


def create_grpc_transport(sandbox: Any, port: int) -> SandboxTransport:
    getter = getattr(sandbox, "get_host", None)
    if not callable(getter):
        raise RuntimeError("sandbox does not support get_host(port)")
    host = str(getter(int(port)) or "").strip()
    if not host:
        raise RuntimeError(f"sandbox get_host returned empty host for port={port}")
    return GrpcSandboxTransport(host)
