src/agent.py

import json
import os
import re
from typing import Dict, Any, Optional, List

from openai import OpenAI

from .environment import TrainingEnv


class AIWorker:
    """Programmatic agent interface — drive the environment step by step."""

    def __init__(self, scenario: str = "log_analysis"):
        self.engine = TrainingEnv(scenario=scenario)
        self.recorder: List[Dict[str, Any]] = []

    def boot(self) -> Dict[str, Any]:
        reset_out = self.engine.reset()
        return {
            "status": "reset",
            "observation": reset_out["observation"],
            "task_name": reset_out["info"]["task_name"],
            "task_instructions": reset_out["info"]["instructions"],
            "max_steps": reset_out["info"]["max_steps"],
            "message": "Environment reset. Read the task instructions and begin solving."
        }

    def invoke(self, action: str, rationale: Optional[str] = None) -> Dict[str, Any]:
        step_out = self.engine.step(action)
        resp = {
            "status": "success" if step_out["info"]["exit_code"] == 0 else "failed",
            "command": action,
            "command_output": step_out["info"].get("command_output", ""),
            "reward": step_out["reward"],
            "task_score": step_out["info"]["task_score"],
            "done": step_out["done"],
            "observation": step_out["observation"],
            "step": step_out["info"]["step"],
            "max_steps": step_out["info"]["max_steps"],
            "task_metadata": step_out["info"]["task_metadata"],
        }
        self.recorder.append({
            "command": action,
            "reasoning": rationale,
            "result": resp,
        })
        if resp["done"]:
            resp["message"] = f"Task completed! Final score: {resp['task_score']:.1f}"
        else:
            resp["message"] = f"Score: {resp['task_score']:.1f} | Steps: {resp['step']}/{resp['max_steps']}"
        return resp

    def context(self) -> str:
        state = self.engine.dump()
        recent = self.recorder[-5:] if self.recorder else []
        recent_cmds = "\n".join(
            f"  [{r['command']}] -> score={r['result']['task_score']}"
            for r in recent
        ) or "No commands executed yet"
        return (
            f"LINUX SRE ENVIRONMENT STATE:\n"
            f"=== CURRENT TASK ===\n"
            f"Name: {self.engine.task.nm}\n"
            f"Difficulty: {self.engine.difficulty}\n"
            f"Description: {self.engine.task.desc}\n\n"
            f"=== PROGRESS ===\n"
            f"Task Score: {self.engine.score:.1f}\n"
            f"Steps Taken: {self.engine.step_count}\n"
            f"Max Steps: {self.engine.limit}\n\n"
            f"=== RECENT COMMANDS ===\n{recent_cmds}\n\n"
            f"=== WORKING DIRECTORY ===\n{self.engine.terminal.cwd}\n"
        )

    def report(self) -> Dict[str, Any]:
        return {
            "task_name": self.engine.task.nm,
            "difficulty": self.engine.difficulty,
            "final_score": self.engine.score,
            "steps_used": self.engine.step_count,
            "max_steps": self.engine.limit,
            "completed": self.engine.finished,
            "efficiency": round(1 - (self.engine.step_count / self.engine.limit), 2)
            if self.engine.finished else 0.0,
            "command_history": self.engine.terminal.history(),
        }


# ======================================================================
#  LLM AGENT (OpenAI client)
# ======================================================================

class LLMAgent:
    """
    LLM-powered agent that autonomously solves scenarios.

    Uses the OpenAI client directly and reads configuration from:
      - API_BASE_URL
      - MODEL_NAME
      - HF_TOKEN

    Usage:
        agent = LLMAgent()
        result = agent.solve("cascading_db_failure")
    """

    def __init__(
        self,
        model: Optional[str] = None,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        max_turns: int = 30,
        verbose: bool = False,
    ):
        self.model = self._normalize_model_name(model or os.environ.get("MODEL_NAME", ""))
        self.api_key = api_key or os.environ.get("HF_TOKEN", "")
        self.base_url = base_url or os.environ.get("API_BASE_URL", "")
        self.max_turns = max_turns
        self.verbose = verbose
        self._client: Optional[OpenAI] = None

    @staticmethod
    def _normalize_model_name(model_name: str) -> str:
        if "/" in model_name:
            provider, candidate = model_name.split("/", 1)
            if provider in {"openai", "anthropic", "gemini"}:
                return candidate
        return model_name

    def _ensure_client(self) -> OpenAI:
        if self._client is None:
            missing = []
            if not self.base_url:
                missing.append("API_BASE_URL")
            if not self.model:
                missing.append("MODEL_NAME")
            if not self.api_key:
                missing.append("HF_TOKEN")
            if missing:
                raise RuntimeError(
                    "Missing required environment variables for LLM agent: "
                    + ", ".join(missing)
                )
            kwargs: Dict[str, Any] = {"api_key": self.api_key, "base_url": self.base_url}
            self._client = OpenAI(**kwargs)
        return self._client

    def _build_messages(self, initial: Dict[str, Any]) -> List[Dict[str, str]]:
        return [
            {
                "role": "system",
                "content": (
                    SystemPrompts.get_sys(initial.get("task_name", "linux-sre"))
                    + "\n\nRespond with a short explanation and a single command."
                    + "\nPreferred format: ```bash\\n<command>\\n```"
                ),
            },
            {
                "role": "user",
                "content": SystemPrompts.format_observation(initial),
            },
        ]

    def solve(self, scenario: str = "log_analysis") -> Dict[str, Any]:
        """Run the agent loop: observe, think, act, and repeat until done."""
        client = self._ensure_client()

        worker = AIWorker(scenario=scenario)
        initial = worker.boot()
        messages = self._build_messages(initial)

        turns = []
        for turn in range(self.max_turns):
            response = client.chat.completions.create(
                model=self.model,
                messages=messages,
                temperature=0.2,
                max_tokens=512,
            )
            assistant_msg = (response.choices[0].message.content or "").strip()
            messages.append({"role": "assistant", "content": assistant_msg})

            command = self._extract_command(assistant_msg)
            if not command:
                messages.append({"role": "user", "content": "Return exactly one shell command."})
                continue

            result = worker.invoke(command, rationale=assistant_msg)
            turns.append({"turn": turn + 1, "command": command, "score": result["task_score"]})

            if self.verbose:
                print(f"[Turn {turn + 1}] $ {command} -> score={result['task_score']:.1f}")

            if result["done"]:
                break

            messages.append(
                {
                    "role": "user",
                    "content": (
                        f"Command output:\n{result['command_output']}\n\n"
                        f"Score: {result['task_score']:.1f} | Step {result['step']}/{result['max_steps']}\n\n"
                        "What command should we run next?"
                    ),
                }
            )

        summary = worker.report()
        summary["turns"] = turns
        summary["model"] = self.model
        summary["base_url"] = self.base_url
        return summary

    @staticmethod
    def extract_command(text: str) -> Optional[str]:
        """Extract a shell command from an LLM response."""
        try:
            payload = json.loads(text)
            if isinstance(payload, dict):
                command = payload.get("command")
                if isinstance(command, str) and command.strip():
                    return command.strip()
        except Exception:
            pass

        fenced = re.search(r"```(?:bash|sh|shell)?\s*\n(.+?)\n```", text, re.DOTALL)
        if fenced:
            for line in fenced.group(1).strip().split("\n"):
                line = line.strip()
                if line and not line.startswith("#"):
                    return line.lstrip("$ ")

        inline = re.search(r"`([^`]+)`", text)
        if inline:
            command = inline.group(1).strip()
            if command:
                return command

        for line in text.split("\n"):
            line = line.strip()
            if line.startswith("$ "):
                return line[2:]
        return None

    @staticmethod
    def _extract_command(text: str) -> Optional[str]:
        return LLMAgent.extract_command(text)


# ======================================================================
#  PROMPT TEMPLATES
# ======================================================================

class SystemPrompts:

    @staticmethod
    def get_sys(task_name: str) -> str:
        return (
            "You are a Linux System Reliability Engineer (SRE) agent. "
            "You are troubleshooting a production incident.\n\n"
            "RULES:\n"
            "- Execute ONE command at a time\n"
            "- Wrap your command in a code block: ```bash\\ncommand\\n```\n"
            "- Briefly explain your reasoning before each command\n"
            "- Work systematically: diagnose first, then fix\n"
            "- Verify your fixes after applying them\n\n"
            "Available commands: cat, grep, ls, chmod, ps, kill, systemctl, "
            "cd, tail, head, find, wc, df, du, top, free, uptime, whoami, "
            "id, hostname, date, touch, mkdir, cp, mv, rm, netstat, curl, "
            "journalctl, crontab, iptables, mount, env, export, sort, uniq, "
            "cut, tr, tee, echo, pwd, dig\n\n"
            "Pipes (|) and redirects (>, >>) are supported.\n\n"
            f"Current task: {task_name}"
        )

    @staticmethod
    def format_observation(initial: Dict[str, Any]) -> str:
        return (
            f"TASK: {initial['task_name']}\n\n"
            f"INSTRUCTIONS:\n{initial['task_instructions']}\n\n"
            f"CURRENT STATE:\n"
            f"- Directory: {initial['observation']['current_directory']}\n"
            f"- Processes:\n{initial['observation']['processes']}\n\n"
            f"You have {initial['max_steps']} steps. Begin troubleshooting."
        )


# ======================================================================
#  DEMO
# ======================================================================

def demo():
    """Quick demo showing the agent solving the medium task."""
    agent = AIWorker(scenario="permission_repair")
    initial = agent.boot()
    print(f"Task: {initial['task_name']}")
    print(f"Scripts: {initial['observation']['filesystem']}")

    commands = [
        ("ls -la /home/user/scripts/cleanup.sh", "Check current permissions"),
        ("chmod 0755 /home/user/scripts/cleanup.sh", "Make script executable"),
        ("ls -la /home/user/scripts/cleanup.sh", "Verify permissions fixed"),
    ]

    for cmd, reason in commands:
        result = agent.invoke(cmd, rationale=reason)
        print(f"\n$ {cmd}  ({reason})")
        print(f"  Output: {result['command_output'][:120]}")
        print(f"  Score: {result['task_score']:.1f} | Done: {result['done']}")
        if result["done"]:
            break

    print(f"\nFinal report: {json.dumps(agent.report(), indent=2)}")


if __name__ == "__main__":
    demo()