inference.py

# pyright: reportMissingImports=false
from __future__ import annotations

import argparse
import asyncio
import json
import os
import re
from dataclasses import dataclass
from itertools import zip_longest
from typing import Any

from openai import OpenAI
import yaml

from env.rewards import RewardCalculator
from inference.metrics import EpisodeMetrics
from inference.model_wrapper import ModelWrapper, score_action_candidate
from inference.prompts import JUDGE_SYSTEM_PROMPT, heuristic_action
from inference.visualize import save_metrics_json, save_reward_curve, save_success_rate_history

try:
    from my_env_v4 import MyEnvV4Action, MyEnvV4Env  # type: ignore[import-not-found]

    EXTERNAL_ENV_AVAILABLE = True
except Exception:
    MyEnvV4Action = None
    MyEnvV4Env = None
    EXTERNAL_ENV_AVAILABLE = False


API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
HF_TOKEN = os.getenv("HF_TOKEN")
LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")

TASK_NAME = os.getenv("MY_ENV_V4_TASK", "cicd-debugger-task")
BENCHMARK = os.getenv("MY_ENV_V4_BENCHMARK", "cicd_debugger_env")

MAX_STEPS_DEFAULT = int(os.getenv("MAX_STEPS", "8"))
TEMPERATURE = float(os.getenv("TEMPERATURE", "0.2"))
MAX_TOKENS = int(os.getenv("MAX_TOKENS", "120"))
OFFLINE_INFERENCE = os.getenv("OFFLINE_INFERENCE", "0") == "1"


DEFAULT_ORIGINAL_CONFIG = """
name: CI
on: [push]
jobs:
  test:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - run: npm ci
      - run: npm tset
""".strip()

DEFAULT_EXPECTED_CONFIG = """
name: CI
on: [push]
jobs:
  test:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - run: npm ci
      - run: npm test
""".strip()

DEFAULT_ERROR_MESSAGE = "command not found"


@dataclass
class LocalObservation:
    config: str
    error_message: str
    logs: str
    last_action_error: str | None = None


@dataclass
class LocalStepResult:
    observation: LocalObservation
    reward: float
    done: bool
    last_action_error: str | None = None


@dataclass
class LocalAction:
    message: str


class LocalCICDDebuggerEnv:
    def __init__(self, original_config: str, expected_config: str, error_message: str):
        self.original_config = original_config
        self.expected_config = expected_config
        self.error_message = error_message
        self.current_config = original_config

    async def reset(self) -> LocalStepResult:
        self.current_config = self.original_config
        obs = LocalObservation(
            config=self.current_config,
            error_message=self.error_message,
            logs="CI failed in test step: npm tset is not a valid command.",
            last_action_error=None,
        )
        return LocalStepResult(observation=obs, reward=0.0, done=False, last_action_error=None)

    async def step(self, action: LocalAction) -> LocalStepResult:
        message = str(action.message or "").strip()
        lower_message = message.lower()
        previous = self.current_config
        step_error: str | None = None
        logs = "No effective change applied."

        if _is_hacking_action(message):
            step_error = "disallowed_hacking_pattern"
            logs = "Rejected unsafe action pattern."
        elif "npm tset" in lower_message and "npm test" in lower_message and "npm tset" in previous:
            self.current_config = previous.replace("npm tset", "npm test")
            logs = "Patched CI command typo from npm tset to npm test."
        elif "replace" in lower_message and "npm test" in lower_message and "npm tset" in previous:
            self.current_config = previous.replace("npm tset", "npm test")
            logs = "Applied replace operation for broken test command."
        elif "npm test" in lower_message and "npm tset" in previous:
            self.current_config = previous.replace("npm tset", "npm test")
            logs = "Applied inferred command fix."

        done = "npm tset" not in self.current_config.lower() and "npm test" in self.current_config.lower()
        reward = 1.0 if done else 0.0
        err_msg = "" if done else self.error_message

        obs = LocalObservation(
            config=self.current_config,
            error_message=err_msg,
            logs=logs,
            last_action_error=step_error,
        )
        return LocalStepResult(observation=obs, reward=reward, done=done, last_action_error=step_error)

    async def close(self) -> None:
        return None


class OpenAIJudgeAdapter:
    def __init__(self, client: OpenAI, model_name: str):
        self.client = client
        self.model_name = model_name

    def evaluate_fix(self, original: str, fixed: str, error: str) -> dict[str, float]:
        prompt = (
            "Evaluate CI config fix quality. Return JSON only with keys correctness, minimalism, quality in [0,1].\n\n"
            f"Original:\n{original}\n\n"
            f"Fixed:\n{fixed}\n\n"
            f"Error:\n{error}\n"
        )
        default = {"correctness": 0.0, "minimalism": 0.0, "quality": 0.0}

        try:
            completion = self.client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {"role": "system", "content": JUDGE_SYSTEM_PROMPT},
                    {"role": "user", "content": prompt},
                ],
                temperature=0.0,
                max_tokens=120,
                stream=False,
            )
            content = (completion.choices[0].message.content or "").strip()
        except Exception:
            return default

        parsed = self._parse_scores(content)
        return parsed if parsed else default

    def _parse_scores(self, content: str) -> dict[str, float] | None:
        decoder = json.JSONDecoder()
        for idx, char in enumerate(content):
            if char != "{":
                continue
            try:
                obj, _ = decoder.raw_decode(content[idx:])
            except json.JSONDecodeError:
                continue
            if isinstance(obj, dict):
                return {
                    "correctness": self._clamp(obj.get("correctness", 0.0)),
                    "minimalism": self._clamp(obj.get("minimalism", 0.0)),
                    "quality": self._clamp(obj.get("quality", 0.0)),
                }

        fallback = {
            "correctness": self._extract_regex(content, "correctness"),
            "minimalism": self._extract_regex(content, "minimalism"),
            "quality": self._extract_regex(content, "quality"),
        }
        if any(value > 0 for value in fallback.values()):
            return fallback
        return None

    def _extract_regex(self, content: str, key: str) -> float:
        match = re.search(rf"{key}\s*[:=\-]\s*([0-9]*\.?[0-9]+)", content, flags=re.IGNORECASE)
        if not match:
            return 0.0
        return self._clamp(match.group(1))

    def _clamp(self, value: Any) -> float:
        try:
            parsed = float(value)
        except (TypeError, ValueError):
            parsed = 0.0
        return max(0.0, min(1.0, parsed))


def log_start(task: str, env_name: str, model: str) -> None:
    print(f"[START] task={_single_line(task)} env={_single_line(env_name)} model={_single_line(model)}", flush=True)


def log_step(step: int, action: str, reward: float, done: bool, error: str | None) -> None:
    done_val = str(done).lower()
    error_val = _single_line(error) if error else "null"
    action_val = _single_line(action)
    print(f"[STEP] step={step} action={action_val} reward={reward:.2f} done={done_val} error={error_val}", flush=True)


def log_end(success: bool, steps: int, rewards: list[float]) -> None:
    rewards_str = ",".join(f"{value:.2f}" for value in rewards)
    print(f"[END] success={str(success).lower()} steps={steps} rewards={rewards_str}", flush=True)


def _single_line(value: Any) -> str:
    return " ".join(str(value).replace("\n", " ").replace("\r", " ").split())


def _safe_float(value: Any) -> float:
    try:
        return float(value or 0.0)
    except (TypeError, ValueError):
        return 0.0


def _extract_observation(result: Any) -> Any:
    return getattr(result, "observation", result)


def _extract_done(result: Any) -> bool:
    return bool(getattr(result, "done", False))


def _extract_reward(result: Any) -> float:
    return _safe_float(getattr(result, "reward", 0.0))


def _extract_error(result: Any, observation: Any) -> str | None:
    result_error = getattr(result, "last_action_error", None)
    if result_error:
        return str(result_error)

    if isinstance(observation, dict):
        obs_err = observation.get("last_action_error")
        if obs_err:
            return str(obs_err)
    else:
        obs_err = getattr(observation, "last_action_error", None)
        if obs_err:
            return str(obs_err)

    return None


def _extract_observation_text(observation: Any) -> str:
    if isinstance(observation, dict):
        for key in ("config", "current_config", "echoed_message"):
            if key in observation and observation.get(key) is not None:
                return str(observation.get(key))
        return str(observation)

    for attr in ("config", "current_config", "echoed_message"):
        if hasattr(observation, attr):
            value = getattr(observation, attr)
            if value is not None:
                return str(value)

    return str(observation)


def _extract_observation_error_message(observation: Any) -> str:
    if isinstance(observation, dict):
        return str(observation.get("error_message", "") or "")
    return str(getattr(observation, "error_message", "") or "")


def _count_changed_lines(previous_text: str, current_text: str) -> int:
    previous_lines = (previous_text or "").splitlines()
    current_lines = (current_text or "").splitlines()
    changed = 0
    for left, right in zip_longest(previous_lines, current_lines, fillvalue=""):
        if left != right:
            changed += 1
    return changed


def _is_hacking_action(action_text: str) -> bool:
    value = (action_text or "").lower()
    patterns = (
        "if: false",
        "when: never",
        "echo \"tests passed\"",
        "echo 'tests passed'",
        "exit 0",
        "force success",
    )
    return any(token in value for token in patterns)


def _classify_action(action_text: str) -> str:
    text = (action_text or "").lower()
    if "read" in text and "log" in text:
        return "read_logs"
    if "analy" in text:
        return "analyze_error"
    if "propose" in text:
        return "propose_fix"
    if any(token in text for token in ("validate", "run test", "pipeline run", "verify")):
        return "validate_fix"
    return "edit_config"


def _select_action(
    model_wrapper: ModelWrapper,
    step: int,
    config_text: str,
    error_message: str,
    history: list[str],
    policy_mode: str,
    trajectories: int,
) -> str:
    mode = (policy_mode or "imp").lower()

    if mode == "sft":
        return heuristic_action(config_text, error_message)

    if mode == "direct":
        return model_wrapper.generate_action(
            step=step,
            config_text=config_text,
            error_message=error_message,
            history=history,
        )

    candidates = model_wrapper.generate_candidates(
        step=step,
        config_text=config_text,
        error_message=error_message,
        history=history,
        count=max(1, int(trajectories)),
    )

    if not candidates:
        return heuristic_action(config_text, error_message)

    observation = f"{config_text}\n{error_message}"
    best = max(candidates, key=lambda item: score_action_candidate(observation, item, _is_hacking_action))
    return best


def _build_action(action_class: Any, message: str) -> Any:
    try:
        return action_class(message=message)
    except TypeError:
        return action_class(message)


async def _load_environment(
    original_config: str,
    expected_config: str,
    error_message: str,
    force_local_env: bool,
) -> tuple[Any, Any]:
    if not force_local_env and EXTERNAL_ENV_AVAILABLE and LOCAL_IMAGE_NAME:
        try:
            env = await MyEnvV4Env.from_docker_image(LOCAL_IMAGE_NAME)
            return env, MyEnvV4Action
        except Exception:
            pass

    env = LocalCICDDebuggerEnv(
        original_config=original_config,
        expected_config=expected_config,
        error_message=error_message,
    )
    return env, LocalAction


def _load_text(raw_value: str | None, file_path: str | None, fallback: str) -> str:
    if raw_value:
        return raw_value
    if file_path:
        with open(file_path, "r", encoding="utf-8") as handle:
            return handle.read().strip()
    return fallback


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run OpenEnv-style CI/CD pipeline debugging inference loop")
    parser.add_argument("--max-steps", type=int, default=MAX_STEPS_DEFAULT)
    parser.add_argument("--task", default=TASK_NAME)
    parser.add_argument("--benchmark", default=BENCHMARK)
    parser.add_argument("--offline", action="store_true", default=OFFLINE_INFERENCE)
    parser.add_argument("--policy-mode", choices=["sft", "imp", "direct"], default="imp")
    parser.add_argument("--trajectories", type=int, default=3)
    parser.add_argument("--force-local-env", action="store_true", default=False)

    parser.add_argument("--original-config", default=None)
    parser.add_argument("--original-config-file", default=None)
    parser.add_argument("--expected-config", default=None)
    parser.add_argument("--expected-config-file", default=None)
    parser.add_argument("--error-message", default=DEFAULT_ERROR_MESSAGE)

    return parser.parse_args()


async def run_episode(args: argparse.Namespace) -> int:
    original_config = _load_text(args.original_config, args.original_config_file, DEFAULT_ORIGINAL_CONFIG)
    expected_config = _load_text(args.expected_config, args.expected_config_file, DEFAULT_EXPECTED_CONFIG)
    error_message = str(args.error_message or DEFAULT_ERROR_MESSAGE)

    env = None
    history: list[str] = []
    steps_taken = 0
    success = False
    metrics = EpisodeMetrics()

    offline_mode = bool(args.offline or not HF_TOKEN)
    client: OpenAI | None = None
    if not offline_mode:
        client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN or "")

    log_start(task=str(args.task), env_name=str(args.benchmark), model=MODEL_NAME)

    try:
        env, action_class = await _load_environment(
            original_config=original_config,
            expected_config=expected_config,
            error_message=error_message,
            force_local_env=bool(args.force_local_env),
        )

        judge_adapter = OpenAIJudgeAdapter(client, MODEL_NAME) if client is not None else None
        reward_calculator = RewardCalculator(llm_judge=judge_adapter)
        model_wrapper = ModelWrapper(
            client=client,
            model_name=MODEL_NAME,
            temperature=TEMPERATURE,
            max_tokens=MAX_TOKENS,
            offline=offline_mode,
        )

        reset_result = await env.reset()
        observation = _extract_observation(reset_result)
        previous_config = original_config
        current_error_message = error_message

        for step in range(1, max(1, int(args.max_steps)) + 1):
            config_text = _extract_observation_text(observation) or previous_config
            obs_error = _extract_observation_error_message(observation)
            if obs_error:
                current_error_message = obs_error

            action_text = _select_action(
                model_wrapper=model_wrapper,
                step=step,
                config_text=config_text,
                error_message=current_error_message,
                history=history,
                policy_mode=str(args.policy_mode),
                trajectories=max(1, int(args.trajectories)),
            )
            action_obj = _build_action(action_class, action_text)

            step_result = await env.step(action_obj)
            observation = _extract_observation(step_result)

            env_reward = _extract_reward(step_result)
            done = _extract_done(step_result)
            step_error = _extract_error(step_result, observation)
            current_config = _extract_observation_text(observation) or config_text
            obs_error = _extract_observation_error_message(observation)
            if obs_error:
                current_error_message = obs_error

            action_type = _classify_action(action_text)
            hacking_attempt = _is_hacking_action(action_text)

            result_for_reward = {
                "previous_config": previous_config,
                "current_config": current_config,
                "fixed_config": current_config,
                "expected_config": expected_config,
                "error": current_error_message,
                "logs_analyzed": "log" in action_text.lower() or action_type == "read_logs",
                "error_diagnosed": action_type in {"analyze_error", "propose_fix", "edit_config", "validate_fix"},
                "fix_proposed": action_type in {"propose_fix", "edit_config"},
                "pipeline_run": action_type == "validate_fix" and step_error is None,
                "tests_passed": done,
                "command_succeeded": step_error is None,
                "changed_files_count": 1 if previous_config != current_config else 0,
                "changed_lines_count": _count_changed_lines(previous_config, current_config),
                "hacking_attempt": hacking_attempt,
            }

            calculated_reward = reward_calculator.calculate_step_reward(
                state={
                    "step_count": step,
                    "previous_config": previous_config,
                    "expected_config": expected_config,
                    "original_config": original_config,
                    "error": current_error_message,
                },
                action=action_type,
                result=result_for_reward,
                original_config=original_config,
                fixed_config=current_config,
                error_message=current_error_message,
                expected_config=expected_config,
                metadata={"broken_token": "npm tset", "fixed_token": "npm test"},
            )

            combined_reward = round(float(calculated_reward) + float(env_reward), 4)
            metrics.add_step(action=action_text, reward=combined_reward, error=step_error, done=done)
            steps_taken = step

            log_step(step=step, action=action_text, reward=combined_reward, done=done, error=step_error)

            history.append(f"step={step} action={_single_line(action_text)} reward={combined_reward:.2f}")
            previous_config = current_config

            if done:
                success = step_error is None and not hacking_attempt
                break

    except Exception:
        success = False
    finally:
        try:
            save_reward_curve(metrics.rewards)
            save_metrics_json(metrics.summary())
            save_success_rate_history([success])
        except Exception:
            pass

        if env is not None:
            try:
                await env.close()
            except Exception:
                pass

        log_end(success=success, steps=steps_taken, rewards=metrics.rewards)

    return 0


def main() -> int:
    args = parse_args()
    return asyncio.run(run_episode(args))


if __name__ == "__main__":
    raise SystemExit(main())