inference.py

from __future__ import annotations

import argparse
import json
import os
import re
import shlex
import sys
import traceback
from typing import Any

import requests
from openai import OpenAI


MAX_STEPS = 40
DEFAULT_ENV_URL = "http://localhost:7860"
ENV_NAME = "secrets_audit"

# Per-difficulty step caps — fewer steps = higher efficiency bonus
STEPS_BY_DIFFICULTY = {
    "easy": 8,     # budget=10, solve in 8 → 0.03 bonus
    "medium": 15,  # budget=20, solve in 15 → 0.0375 bonus
    "hard": 25,    # budget=30, solve in 25 → 0.025 bonus
}

TASK_DIFFICULTY = {
    1: "easy", 2: "easy", 3: "easy", 4: "easy", 5: "easy",
    6: "medium", 7: "medium", 8: "medium", 9: "medium", 10: "medium",
    11: "hard", 12: "hard", 13: "hard",
}
PRIMARY_FILE_BY_TASK_ID = {
    1: "config.py",
    2: "db.py",
    3: "settings.js",
    4: "logger.py",
    5: ".env",
    6: "utils.py",
    7: "deploy.yml",
    8: "app.toml",
    9: "migrate.sql",
    10: "deploy.sh",
    11: "service_a.py",
    12: "crypto.py",
    13: "config.py",
}


def normalize_action(raw: str) -> str:
    text = raw.strip()
    if not text:
        return "true"
    if _looks_like_provider_error(text):
        return "true"

    # Try extracting from ```bash ... ``` fences first
    fence_match = re.search(r"```(?:bash|sh)?\s*(.*?)```", text, re.DOTALL)
    if fence_match:
        text = fence_match.group(1).strip()
    if _looks_like_provider_error(text):
        return "true"

    # Try extracting from XML tool_call (Minimax, etc.)
    xml_match = re.search(r'<parameter\s+name="command">(.*?)</parameter>', text, re.DOTALL)
    if xml_match:
        text = xml_match.group(1).strip()

    lines = [
        line.strip()
        for line in text.splitlines()
        if line.strip() and not line.strip().startswith(("#", "-", "*"))
    ]
    if not lines:
        return "true"

    # If first line looks like a shell command, use it
    first_line = lines[0]
    if _looks_like_shell(first_line):
        return first_line

    # If first line is prose, scan remaining lines for a shell command
    for line in lines[1:]:
        if _looks_like_shell(line):
            return line

    return "true"


def _looks_like_shell(line: str) -> bool:
    if not line:
        return False
    prose_prefixes = (
        "here",
        "this",
        "i ",
        "i'",
        "you ",
        "the ",
        "to ",
        "we ",
        "run ",
        "use ",
        "first ",
        "let ",
        "let's",
        "looking",
        "now ",
        "next ",
        "since ",
        "note",
        "okay",
        "sure",
        "great",
        "step ",
        "<",
    )
    lowered = line.lower()
    if lowered.startswith(prose_prefixes):
        return False
    if _looks_like_provider_error(line):
        return False
    return True


def _looks_like_provider_error(text: str) -> bool:
    lowered = text.strip().lower()
    error_markers = (
        "internal server error",
        "server error",
        "bad gateway",
        "gateway timeout",
        "service unavailable",
        "rate limit",
        "too many requests",
        "upstream error",
        "provider error",
        "api error",
        "error code:",
        "request failed",
    )
    return any(lowered.startswith(marker) for marker in error_markers)


def stderr_log(message: str) -> None:
    print(message, file=sys.stderr, flush=True)


def format_field(value: Any) -> str:
    text = "none" if value is None else re.sub(r"\s+", " ", str(value).strip())
    if not text:
        text = "none"
    if " " in text:
        return json.dumps(text)
    return text


def bool_text(value: bool) -> str:
    return "true" if value else "false"


def stdout_tag(tag: str, **fields: Any) -> None:
    parts = [f"[{tag}]"]
    for key, value in fields.items():
        parts.append(f"{key}={format_field(value)}")
    print(" ".join(parts), flush=True)


def build_prompt(
    state: dict[str, Any],
    recent_actions: list[str] | None = None,
    stuck_warning: str | None = None,
) -> str:
    sanitized_state = sanitize_state_for_prompt(state)
    session = sanitized_state["session"]
    last_result = session.get("last_result") or {}
    recent_actions = recent_actions or []
    recent_summary = ", ".join(recent_actions[-5:]) if recent_actions else "none"
    health_stdout = sanitize_observation_text(session.get("health_stdout", ""))
    health_stderr = sanitize_observation_text(session.get("health_stderr", ""))
    observation = sanitize_observation_text(session.get("observation", ""))
    stuck_section = f"\n[CRITICAL WARNING: YOU ARE STUCK] {stuck_warning}\n" if stuck_warning else ""

    task_id = session.get("task_id", 0)
    primary_file = PRIMARY_FILE_BY_TASK_ID.get(task_id if isinstance(task_id, int) else int(str(task_id).replace('task_', '') or 0), "")
    steps_taken = session.get("steps_taken", 0)
    reward = session.get("reward", 0)

    # Task-specific hints
    task_hints = ""
    if task_id in (13, "13"):
        task_hints = (
            "\n[TASK 13 — GIT HISTORY LEAK]\n"
            "The secret is in git history. Deleting files will NOT work.\n"
            "You MUST use: git filter-repo --replace-text <(echo 'SECRET_VALUE==>REDACTED') --force\n"
            "First run: git log --all -p | grep -i 'key\\|token\\|secret\\|password' to find the secret value.\n"
        )
    elif reward > 0.4 and reward < 0.95:
        task_hints = (
            "\n[PARTIAL FIX DETECTED — CHECK GIT HISTORY]\n"
            "Reward is stuck around 0.5. The working file is fixed but git log still has the old secret.\n"
            "Run: git filter-repo --replace-text <(echo 'OLD_SECRET_VALUE==>REDACTED') --force\n"
            "Or: git filter-branch --tree-filter 'sed -i s/OLD_SECRET/REDACTED/g FILE' -- --all\n"
        )

    # Efficiency guidance
    efficiency_hint = ""
    if steps_taken == 0:
        efficiency_hint = f"\nPriority: Read {primary_file} first, then fix the secret in one sed command.\n" if primary_file else ""
    elif steps_taken >= 5 and reward < 0.1:
        efficiency_hint = "\n[URGENT] 5+ steps with no progress. Stop reading files. Fix the secret NOW with sed.\n"

    return (
        "You are a security engineer remediating leaked secrets in a benchmark workspace.\n"
        f"Task {session['task_id']}: {session['title']}\n"
        f"Description: {session['description']}\n"
        f"Workspace: {session['workspace']}\n\n"
        f"--- STATUS ---\n"
        f"Reward: {reward} | Leaks remaining: {session['current_leaks']} | Health: {session['health_score']}\n"
        f"Steps taken: {steps_taken}\n"
        f"Recent actions: {recent_summary}\n\n"
        f"--- PREVIOUS RESULT ---\n"
        f"Action: {last_result.get('action', 'none')}\n"
        f"Stdout: {last_result.get('stdout', '')[:500]}\n"
        f"Stderr: {last_result.get('stderr', '')[:300]}\n\n"
        f"--- OBSERVATION ---\n{observation[:800]}\n\n"
        f"--- HEALTH ---\n{health_stdout[:300]}\n{health_stderr[:200]}\n\n"
        "=== RULES (MUST FOLLOW) ===\n"
        "1. Return EXACTLY ONE bash command. No explanation, no markdown, no chaining (&&, ;, |).\n"
        "2. Use relative paths only. You are inside the task workspace.\n"
        "3. When replacing secrets, use os.getenv('VAR_NAME') and add 'import os' at the top of the file.\n"
        "4. Do NOT repeat the same action. Each step must make progress.\n"
        "5. After fixing a file, run: gitleaks detect --no-git --source . -v  to verify.\n"
        "6. If reward is ~0.5, the file is fixed but git history still leaks. Use git filter-repo.\n"
        "7. If health drops to 0, your edit broke the code. Fix the syntax error immediately.\n"
        f"{task_hints}"
        f"{efficiency_hint}"
        f"{stuck_section}"
    )


def call_model(client: OpenAI, model_name: str, prompt: str) -> str:
    response = client.chat.completions.create(
        model=model_name,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.1,
    )
    content = extract_response_text(response)
    if content:
        return content

    stderr_log(f"model_response_without_text={safe_model_dump(response)}")
    return ""


def extract_response_text(response: Any) -> str:
    choices = getattr(response, "choices", None)
    if choices:
        first_choice = choices[0]
        message = getattr(first_choice, "message", None)
        if message is not None:
            content = getattr(message, "content", None)
            extracted = coerce_content_to_text(content)
            if extracted:
                return extracted

    output_text = getattr(response, "output_text", None)
    if output_text:
        return str(output_text)

    output = getattr(response, "output", None)
    if output:
        extracted = coerce_content_to_text(output)
        if extracted:
            return extracted

    dumped = safe_model_dump(response)
    return extract_text_from_dump(dumped)


def coerce_content_to_text(content: Any) -> str:
    if content is None:
        return ""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts: list[str] = []
        for item in content:
            if isinstance(item, str):
                parts.append(item)
                continue
            if isinstance(item, dict):
                text = item.get("text")
                if isinstance(text, str):
                    parts.append(text)
                    continue
                inner = item.get("content")
                if isinstance(inner, str):
                    parts.append(inner)
            else:
                text = getattr(item, "text", None)
                if isinstance(text, str):
                    parts.append(text)
        return "\n".join(part for part in parts if part).strip()
    if isinstance(content, dict):
        for key in ("text", "content", "output_text"):
            value = content.get(key)
            if isinstance(value, str) and value.strip():
                return value
        return ""
    return str(content).strip()


def safe_model_dump(response: Any) -> dict[str, Any]:
    if hasattr(response, "model_dump"):
        try:
            dumped = response.model_dump()
            if isinstance(dumped, dict):
                return dumped
        except Exception:
            return {"repr": repr(response)}
    if isinstance(response, dict):
        return response
    return {"repr": repr(response)}


def extract_text_from_dump(payload: Any) -> str:
    if isinstance(payload, str):
        return payload.strip()
    if isinstance(payload, list):
        for item in payload:
            extracted = extract_text_from_dump(item)
            if extracted:
                return extracted
        return ""
    if isinstance(payload, dict):
        for key in ("content", "text", "output_text"):
            value = payload.get(key)
            if isinstance(value, str) and value.strip():
                return value.strip()
            extracted = extract_text_from_dump(value)
            if extracted:
                return extracted
        for value in payload.values():
            extracted = extract_text_from_dump(value)
            if extracted:
                return extracted
    return ""


def post_json(base_url: str, path: str, payload: dict[str, Any], timeout: int) -> dict[str, Any]:
    response = requests.post(f"{base_url}{path}", json=payload, timeout=timeout)
    response.raise_for_status()
    return response.json()


def is_done(state: dict[str, Any]) -> bool:
    session = state["session"]
    return float(session["reward"]) >= 0.99


def extract_error(session: dict[str, Any]) -> str:
    session_error = (session.get("error") or "").strip()
    if session_error and session_error != "none":
        return session_error
    last_result = session.get("last_result") or {}
    if last_result.get("timed_out"):
        return f"timeout:{last_result.get('stderr') or 'command timed out'}"
    if int(last_result.get("exit_code", 0)) != 0:
        return (last_result.get("stderr") or last_result.get("stdout") or "command failed").strip()
    return "none"


def parse_task_id(task_value: str) -> int:
    text = str(task_value).strip()
    match = re.fullmatch(r"task_(\d+)", text)
    if match:
        return int(match.group(1))
    return int(text)


def detect_repeated_action(actions: list[str], rewards: list[float]) -> tuple[str, str] | None:
    if len(actions) < 3 or len(rewards) < 3:
        return None
    if actions[-1] == actions[-2] == actions[-3] and rewards[-1] == rewards[-2] == rewards[-3]:
        warning = (
            f"The last three steps repeated {actions[-1]!r} and reward stayed at {rewards[-1]}. "
            "You are stuck. Choose a different single command that inspects or patches the primary source file."
        )
        return actions[-1], warning
    return None


def choose_fallback_action(task_id: int, repeated_action: str) -> str:
    candidate = PRIMARY_FILE_BY_TASK_ID.get(task_id)
    if "git status" in repeated_action:
        if candidate:
            return f"sed -n '1,200p' {shlex.quote(candidate)}"
        return "find . -maxdepth 2 -type f"
    if repeated_action.startswith("cat ") or repeated_action.startswith("sed -n"):
        return "pytest -q"
    if candidate:
        return f"sed -n '1,200p' {shlex.quote(candidate)}"
    return "find . -maxdepth 2 -type f"


def sanitize_state_for_prompt(state: dict[str, Any]) -> dict[str, Any]:
    return sanitize_observation_value(state)


def sanitize_observation_value(value: Any) -> Any:
    if isinstance(value, dict):
        return {key: sanitize_observation_value(item) for key, item in value.items()}
    if isinstance(value, list):
        return [sanitize_observation_value(item) for item in value]
    if isinstance(value, str):
        return sanitize_observation_text(value)
    return value


def sanitize_observation_text(text: str) -> str:
    if not text:
        return text
    sanitized = text
    repo_root = os.getcwd()
    sanitized = sanitized.replace(f"{repo_root}/", "")
    sanitized = sanitized.replace(repo_root, ".")
    sanitized = re.sub(r"/[^/\s]*/runtime/session_[^/\s]+/", "", sanitized)
    sanitized = re.sub(r"/home/[^/\s]+/", "", sanitized)
    sanitized = re.sub(r"\.{2,}", ".", sanitized)
    return sanitized


def enforce_atomic_action(action: str) -> str:
    """Reject multi-command chaining but allow ${VAR} and curly braces in sed/grep."""
    if not action:
        return "true"
    # Reject && || ; chaining
    if re.search(r"&&|\|\|", action):
        return "true"
    # Reject semicolons that aren't inside quotes
    # Simple heuristic: if there's a ; outside of quotes, reject
    stripped = re.sub(r"'[^']*'|\"[^\"]*\"", "", action)  # remove quoted strings
    if ";" in stripped:
        return "true"
    return action


def run_single_task(
    client: OpenAI,
    model_name: str,
    env_url: str,
    task_id: int,
    task_label: str,
) -> float:
    """Run a single task episode and return the graded score."""
    rewards: list[float] = []
    step_infos: list[dict] = []
    actions: list[str] = []
    success = False
    steps_run = 0
    final_score = 0.0

    # Must use [START] tag — validator counts [START]/[END] pairs per task
    print(f"[START] task={task_label} env={ENV_NAME} model={model_name}", flush=True)

    # Determine step limit based on difficulty
    difficulty = TASK_DIFFICULTY.get(task_id, "medium")
    max_steps = STEPS_BY_DIFFICULTY.get(difficulty, MAX_STEPS)
    stderr_log(f"task={task_label} difficulty={difficulty} max_steps={max_steps}")

    try:
        state = post_json(env_url, "/reset", {"task_id": task_id}, timeout=30)
        final_score = float(state["session"]["reward"])

        # Auto-inject first action: read the primary file (saves 1-2 LLM calls)
        primary_file = PRIMARY_FILE_BY_TASK_ID.get(task_id)
        if primary_file:
            auto_action = f"cat {primary_file}"
            stderr_log(f"task={task_label} auto_action={auto_action}")
            state = post_json(env_url, "/step", {"action": auto_action}, timeout=90)
            session = state["session"]
            final_score = float(session["reward"])
            actions.append(auto_action)
            rewards.append(final_score)
            step_infos.append({"reward": final_score, "action": auto_action, "step": 0})
            done = is_done(state)
            action_trunc = auto_action[:200].replace("\n", " ")
            print(f"[STEP] step=0 action={action_trunc} reward={final_score:.2f} done={str(done).lower()} error=null", flush=True)
            steps_run = 1
            if done:
                success = True

        if not success:
            for step_num in range(1, max_steps + 1):
                stderr_log(f"task={task_label} step={step_num} building prompt")
                repeated = detect_repeated_action(actions, rewards)
                forbidden_action = None
                stuck_warning = None
                if repeated:
                    forbidden_action, stuck_warning = repeated
                    stderr_log(f"task={task_label} step={step_num} repeated_action_detected={forbidden_action!r}")
                prompt = build_prompt(state, recent_actions=actions, stuck_warning=stuck_warning)
                raw = call_model(client, model_name, prompt)
                stderr_log(f"task={task_label} step={step_num} raw_response={raw!r}")

                action = normalize_action(raw)
                atomic_action = enforce_atomic_action(action)
                if atomic_action != action:
                    stderr_log(f"task={task_label} step={step_num} rejected_non_atomic_action={action!r}")
                    action = atomic_action
                if forbidden_action and action == forbidden_action:
                    fallback_action = choose_fallback_action(task_id, forbidden_action)
                    stderr_log(
                        f"task={task_label} step={step_num} overriding_repeated_action={forbidden_action!r} fallback={fallback_action!r}"
                    )
                    action = fallback_action
                state = post_json(env_url, "/step", {"action": action}, timeout=90)

                session = state["session"]
                final_score = float(session["reward"])
                actions.append(action)
                rewards.append(final_score)

                # Collect step info for /grader call
                step_infos.append({
                    "reward": final_score,
                    "action": action,
                    "step": step_num,
                })

                done = is_done(state)
                error_text = extract_error(session)

                # Exact format: [STEP] step=N action=... reward=0.50 done=false error=null
                action_trunc = action[:200].replace("\n", " ")
                done_val = str(done).lower()
                error_val = error_text if error_text and error_text != "none" else "null"
                print(f"[STEP] step={step_num} action={action_trunc} reward={final_score:.2f} done={done_val} error={error_val}", flush=True)

                steps_run = step_num
                if done:
                    success = True
                    break

    except Exception as exc:
        stderr_log(f"task={task_label} fatal_error={exc!r}")
        traceback.print_exc(file=sys.stderr)

    # Call POST /grader to get the official score (matching reference project)
    try:
        grade_result = post_json(env_url, "/grader", {
            "task_id": task_label,
            "step_rewards": rewards,
            "step_infos": step_infos,
        }, timeout=30)
        graded_score = float(grade_result.get("score", final_score))
    except Exception as exc:
        stderr_log(f"task={task_label} grader_call_failed={exc!r}")
        graded_score = final_score

    # Clamp to strict (0, 1) for validator
    graded_score = max(0.01, min(0.99, graded_score))

    # Must use [END] tag — validator counts [START]/[END] pairs
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(f"[END] success={str(success).lower()} steps={steps_run} score={graded_score:.3f} rewards={rewards_str}", flush=True)

    return graded_score


# Default tasks to run — validator requires at least 3
DEFAULT_TASKS = "task_1,task_2,task_3"


def main() -> None:
    parser = argparse.ArgumentParser(description="Meta OpenEnv Round 1 inference loop.")
    parser.add_argument(
        "--task-id",
        default=os.environ.get("TASK_ID", DEFAULT_TASKS),
        help="Comma-separated list of task IDs to run (e.g. task_1,task_2,task_3)",
    )
    args = parser.parse_args()

    api_base_url = os.environ.get("API_BASE_URL", "https://openrouter.ai/api/v1").rstrip("/")
    hf_token = os.environ.get("HF_TOKEN", "")
    model_name = os.environ.get("MODEL_NAME", "Qwen/Qwen2.5-72b-instruct")
    env_url = os.environ.get("ENV_URL", DEFAULT_ENV_URL).rstrip("/")

    client = OpenAI(base_url=api_base_url, api_key=hf_token, timeout=60)

    # Parse comma-separated task list
    raw_tasks = str(args.task_id).strip()
    if "," in raw_tasks:
        task_labels = [t.strip() for t in raw_tasks.split(",") if t.strip()]
    else:
        task_labels = [raw_tasks]

    # Ensure at least 3 tasks for the validator
    if len(task_labels) < 3:
        all_defaults = ["task_1", "task_2", "task_3"]
        for t in all_defaults:
            if t not in task_labels:
                task_labels.append(t)
            if len(task_labels) >= 3:
                break

    # No global [START]/[END] — each task emits its own [START]/[END] pair
    # The validator counts how many [END] lines have valid scores
    all_scores: dict[str, float] = {}

    for task_label in task_labels:
        task_id = parse_task_id(task_label)
        score = run_single_task(client, model_name, env_url, task_id, task_label)
        all_scores[task_label] = score

    # Summary to stderr only (not parsed by validator)
    stderr_log(f"tasks_run={len(all_scores)} avg_score={round(sum(all_scores.values()) / max(len(all_scores), 1), 4)}")


if __name__ == "__main__":
    main()