#!/usr/bin/env python3
"""
inference.py — Sentinel Env Multi-Agent Inference Script
OpenEnv 2026 Architecture — Session-isolated, WebSocket-ready
Follows [START]/[STEP]/[END] stdout format.
Runs ONE task per invocation via TASK_NAME env var.
"""

import os
import json
import requests
from openai import OpenAI

# ── Credentials ──
API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
MODEL_NAME   = os.getenv("MODEL_NAME",   "meta-llama/Meta-Llama-3-8B-Instruct")
HF_TOKEN     = os.getenv("HF_TOKEN")

if HF_TOKEN is None:
    raise ValueError("HF_TOKEN environment variable is required")

SPACE_URL = os.getenv("SPACE_URL", "https://rudrapatel-1908-sentinel-env.hf.space")
TASK_NAME = os.getenv("TASK_NAME", "easy-lockdown")
BENCHMARK = "sentinel_env"
MAX_STEPS = 8

client = OpenAI(api_key=HF_TOKEN, base_url=API_BASE_URL)

# ── HTTP Helpers ──
def env_reset(task_id: str) -> dict:
    r = requests.post(f"{SPACE_URL}/reset", json={"task_id": task_id}, timeout=30)
    r.raise_for_status()
    return r.json()

def env_step(command: str, target_id: str) -> dict:
    r = requests.post(f"{SPACE_URL}/step",
                      json={"command": command, "target_id": target_id}, timeout=30)
    r.raise_for_status()
    return r.json()

def warroom_reset() -> dict:
    r = requests.post(f"{SPACE_URL}/warroom/reset",
                      json={"task_id": "red-vs-blue"}, timeout=30)
    r.raise_for_status()
    return r.json()

def warroom_step(att_type, att_res, scn_type, scn_tip, scn_res, rem_type, rem_tip, rem_res) -> dict:
    payload = {
        "attacker":   {"action_type": att_type, "target_resource": att_res},
        "scanner":    {"action_type": scn_type, "target_ip": scn_tip, "target_resource": scn_res},
        "remediator": {"action_type": rem_type, "target_ip": rem_tip, "target_resource": rem_res},
    }
    r = requests.post(f"{SPACE_URL}/warroom/step", json=payload, timeout=30)
    r.raise_for_status()
    return r.json()

# ── LLM: Single-agent action ──
def get_single_action(task_id: str, obs: dict) -> dict:
    terminal  = obs.get("terminal_output", "")
    inventory = json.dumps(obs.get("inventory", []), indent=2)

    rules = {
        "easy-lockdown":       'command="lockdown", target_id="s3-vault"',
        "easy-secrets":        'Step 1: command="audit", target_id="api-key-01" → Step 2: command="revoke", target_id="api-key-01" → Step 3: command="rotate", target_id="api-key-01"',
        "medium-access":       'command="revoke_admin", target_id="user-dev-01"',
        "medium-mfa":          'Step 1: command="audit", target_id="iam-users" → Then: command="enforce_mfa", target_id="admin-alice" (repeat for admin-bob and admin-carol)',
        "hard-breach":         'Step 1: command="block_ip", target_id="attacker-ip" → Step 2: command="close_port", target_id="web-server"',
        "critical-ransomware": 'Step 1: command="isolate", target_id="db-server" → Step 2: command="revoke_sessions", target_id="active-sessions" → Step 3: command="restore_backup", target_id="db-server"',
        "expert-apt":          'Step 1: command="detect_c2", target_id="c2-beacon" → Step 2: command="block_outbound", target_id="outbound-fw" → Step 3: command="isolate_host", target_id="infected-host" → Step 4: command="patch_vulnerability", target_id="vuln-cve-2024"',
    }

    prompt = f"""You are a cybersecurity AI agent for Sentinel-Env.

Task: {task_id}
Terminal output: {terminal}
Current inventory: {inventory}

Correct action sequence for this task:
{rules.get(task_id, 'Follow the terminal instructions.')}

Based on the terminal output, choose the NEXT correct action.
Respond ONLY with a JSON object. No markdown. No explanation.
Example: {{"command": "lockdown", "target_id": "s3-vault"}}"""

    resp = client.chat.completions.create(
        model=MODEL_NAME,
        messages=[{"role": "user", "content": prompt}],
        max_tokens=64,
        temperature=0,
    )
    raw = resp.choices[0].message.content.strip()
    raw = raw.replace("```json", "").replace("```", "").strip()
    return json.loads(raw)

# ── LLM: Multi-agent actions ──
def _ask_llm(prompt: str) -> dict:
    try:
        resp = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[{"role": "user", "content": prompt}],
            max_tokens=80, temperature=0,
        )
        raw = resp.choices[0].message.content.strip().replace("```json","").replace("```","").strip()
        return json.loads(raw)
    except Exception:
        return {}

def get_multi_actions(obs: dict) -> dict:
    scn_obs = obs.get("scanner_obs", {})
    rem_obs = obs.get("remediator_obs", {})
    att_obs = obs.get("attacker_obs", {})

    att = _ask_llm(f"""You are the Red Team attacker.
Visible resources: {att_obs.get('visible_resources',[])}
Position: {att_obs.get('current_position','unknown')}
Data stolen: {att_obs.get('data_exfiltrated',False)}

action_type options: port_scan, privilege_escalation, lateral_movement, data_exfiltration, hide_tracks
target_resource options: s3-vault, db-server, secret-vault, web-server, user-admin

Respond ONLY with JSON: {{"action_type":"...","target_resource":"..."}}""")

    scn = _ask_llm(f"""You are the Blue Team Scanner.
IP connections: {scn_obs.get('ip_connections',[])}
Flagged threats: {scn_obs.get('flagged_threats',[])}
Logs: {scn_obs.get('security_logs',[])[-2:]}

action_type options: scan_logs, flag_threat, share_intel
Respond ONLY with JSON: {{"action_type":"...","target_ip":"45.33.32.156","target_resource":"web-server"}}""")

    rem = _ask_llm(f"""You are the Blue Team Remediator.
Open ports: {rem_obs.get('open_ports',[])}
Locked resources: {rem_obs.get('locked_resources',[])}
Intel from Scanner: {rem_obs.get('received_intel',[])}

action_type options: block_ip, close_port, lockdown, revoke_access
Respond ONLY with JSON: {{"action_type":"...","target_ip":"45.33.32.156","target_resource":"web-server"}}""")

    # Safe defaults
    if not att.get("action_type"): att = {"action_type": "port_scan",  "target_resource": "web-server"}
    if not scn.get("action_type"): scn = {"action_type": "scan_logs",  "target_resource": "web-server"}
    if not rem.get("action_type"): rem = {"action_type": "close_port", "target_resource": "web-server"}

    return {"attacker": att, "scanner": scn, "remediator": rem}

# ── Single-agent runner ──
def run_single_task(task_id: str) -> None:
    result  = env_reset(task_id)
    obs     = result["observation"]
    rewards = []
    done    = False
    success = False

    print(f"[START] task={task_id} env={BENCHMARK} model={MODEL_NAME}", flush=True)

    for step_num in range(1, MAX_STEPS + 1):
        action_str = "null"
        reward     = 0.05
        error      = "null"
        try:
            action     = get_single_action(task_id, obs)
            cmd        = action["command"]
            tgt        = action["target_id"]
            action_str = f"{cmd}({tgt})"
            result     = env_step(cmd, tgt)
            obs        = result["observation"]
            reward     = float(result["reward"])
            done       = result["done"]
        except Exception as e:
            error = str(e).replace("\n", " ")

        rewards.append(reward)
        print(f"[STEP] step={step_num} action={action_str} "
              f"reward={reward:.2f} done={str(done).lower()} error={error}", flush=True)
        if done:
            success = True
            break

    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(f"[END] success={str(success).lower()} steps={len(rewards)} "
          f"rewards={rewards_str}", flush=True)

# ── Multi-agent runner ──
def run_warroom_task() -> None:
    result  = warroom_reset()
    obs     = result
    rewards = []
    done    = False
    success = False

    print(f"[START] task=red-vs-blue env={BENCHMARK} model={MODEL_NAME}", flush=True)

    for step_num in range(1, MAX_STEPS + 1):
        action_str = "null"
        reward     = 0.05
        error      = "null"
        try:
            acts = get_multi_actions(obs)
            att  = acts["attacker"]
            scn  = acts["scanner"]
            rem  = acts["remediator"]

            action_str = (
                f"red:{att.get('action_type','?')}({att.get('target_resource','?')})|"
                f"scan:{scn.get('action_type','?')}|"
                f"rem:{rem.get('action_type','?')}"
            )

            result  = warroom_step(
                att.get("action_type","port_scan"),   att.get("target_resource","web-server"),
                scn.get("action_type","scan_logs"),   scn.get("target_ip",""),  scn.get("target_resource","web-server"),
                rem.get("action_type","close_port"),  rem.get("target_ip",""),  rem.get("target_resource","web-server"),
            )
            obs     = result.get("observation", obs)
            rwds    = result.get("rewards", {})
            done    = result.get("done", False)
            blue    = (rwds.get("scanner", 0.05) + rwds.get("remediator", 0.05)) / 2
            reward  = round(max(0.01, min(0.99, blue)), 2)

        except Exception as e:
            error = str(e).replace("\n", " ")

        rewards.append(reward)
        print(f"[STEP] step={step_num} action={action_str} "
              f"reward={reward:.2f} done={str(done).lower()} error={error}", flush=True)
        if done:
            att_obs = obs.get("attacker_obs", {}) if isinstance(obs, dict) else {}
            success = not att_obs.get("data_exfiltrated", True)
            break

    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(f"[END] success={str(success).lower()} steps={len(rewards)} "
          f"rewards={rewards_str}", flush=True)

# ── Main ──
# ── Main — runs ONLY the task in TASK_NAME ──
def main():
    try:
        requests.get(f"{SPACE_URL}/health", timeout=15).raise_for_status()
    except Exception:
        print("[END] success=false steps=0 rewards=0.05", flush=True)
        raise SystemExit(1)

    valid_tasks = [
        "easy-lockdown",
        "easy-secrets",
        "medium-access",
        "medium-mfa",
        "hard-breach",
        "critical-ransomware",
        "expert-apt",
        "red-vs-blue",
    ]

    if TASK_NAME not in valid_tasks:
        print("[END] success=false steps=0 rewards=0.05", flush=True)
        raise SystemExit(1)

    try:
        if TASK_NAME == "red-vs-blue":
            run_warroom_task()
        else:
            run_single_task(TASK_NAME)
    except Exception:
        print("[END] success=false steps=0 rewards=0.05", flush=True)
        raise SystemExit(1)

if __name__ == "__main__":
    main()