#!/usr/bin/env python3
"""
inference.py — LLM-powered SRE Agent for ResilientAgent-Prod.

Uses the OpenAI-compatible API (Groq / HuggingFace / etc.) to diagnose and
resolve ML production incidents. Reads API_BASE_URL, MODEL_NAME, HF_TOKEN
from the environment (hackathon grader injects these automatically).

Outputs structured logs in [START] [STEP] [END] format for hackathon evaluation.
"""

import os
import sys
import json
from typing import Optional, List

# Ensure local imports work
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from dotenv import load_dotenv
from openai import OpenAI

from server.resilientagent_prod_environment import ResilientAgentEnvironment
from models import ResilientAgentAction

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
load_dotenv()

API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
MODEL_NAME   = os.getenv("MODEL_NAME", "gpt-4")
HF_TOKEN     = os.getenv("HF_TOKEN")
OPENAI_API   = os.getenv("OPENAI_API_KEY")

api_key = HF_TOKEN if HF_TOKEN else OPENAI_API

# Initialize OpenAI client
if not api_key or api_key.strip() == "":
    print("ERROR: Neither HF_TOKEN nor OPENAI_API_KEY environment variable is set!", file=sys.stderr)
    print("Evaluator must inject: API_BASE_URL, MODEL_NAME, and an API Key", file=sys.stderr)
    sys.exit(1)

try:
    client = OpenAI(base_url=API_BASE_URL, api_key=api_key)
except Exception as e:
    print(f"ERROR: Failed to initialize OpenAI client: {e}", file=sys.stderr)
    sys.exit(1)

# ---------------------------------------------------------------------------
# System prompt — gives the LLM full context about the environment
# ---------------------------------------------------------------------------
SYSTEM_PROMPT = """\
You are an autonomous SRE agent that diagnoses and resolves ML production incidents.

## Available actions (pick exactly ONE per step)
check_metrics, read_logs, check_deployment, analyze_drift,
scale_service, rollback_model, optimize_batch, restart_service,
verify_fix, notify_team

## Available targets
inference_service, ml_model, primary_model, fallback_model

## Critical rules
1. NEVER repeat the same (action, target) pair you already used.
2. Follow this general pattern: diagnose first → apply a fix → verify_fix.
3. Reply ONLY with a JSON object:  {"action_type": "...", "target": "..."}
   No markdown fences, no extra text.

## Task-Specific Guidance

### Latency Spike (latency_p99 > 1000ms, gpu_memory_exhaustion in logs)
CORRECT SEQUENCE: check_metrics → read_logs → optimize_batch(inference_service) → verify_fix
- The fix is optimize_batch to reduce GPU memory pressure
- Target: inference_service

### Prediction Drift (accuracy < 0.8, drift_score > 0.5, schema mismatch in logs)
CORRECT SEQUENCE: analyze_drift(ml_model) → check_deployment(ml_model) → rollback_model(ml_model) → verify_fix
- Root cause: data_pipeline_schema_change
- The fix is rollback_model to previous version
- Target: ml_model (NOT inference_service)

### Cascading Failure (primary_model down, fallback_model degraded, OOM in logs)
CORRECT SEQUENCE: check_metrics → read_logs → restart_service(primary_model) → scale_service(fallback_model) → verify_fix
- First fix: restart_service on primary_model (to recover from OOM)
- Second fix: scale_service on fallback_model (to handle load)
- Targets: primary_model and fallback_model (NOT inference_service)

## Action Meanings
- analyze_drift: Check for model/data drift issues
- rollback_model: Revert to previous model version (use for schema/pipeline issues)
- optimize_batch: Reduce batch size to fix GPU memory issues
- restart_service: Restart crashed/dead service (use for OOM/memory leaks)
- scale_service: Add capacity to handle load (use for degraded/high-load services)
"""


def build_user_prompt(task_id: str, obs, history: list[dict]) -> str:
    """Build a rich user prompt with observation + history."""
    obs_summary = {
        "task_id": task_id,
        "alert_status": obs.alert_status,
        "metrics": obs.metrics,
        "recent_logs": obs.recent_logs[:3],
    }

    history_str = ""
    if history:
        history_str = "\n\nActions already taken (DO NOT repeat these):\n"
        for i, h in enumerate(history, 1):
            history_str += f"  {i}. {h['action_type']} -> {h['target']}  (reward={h['reward']:.3f})\n"

    return (
        f"Current observation:\n{json.dumps(obs_summary, indent=2)}"
        f"{history_str}"
        f"\n\nWhat is your next action?"
    )


def get_llm_action(task_id: str, obs, history: list[dict]) -> tuple[dict, Optional[str]]:
    """Ask the LLM for the next action. Returns (action_dict, error_string or None)."""
    prompt = build_user_prompt(task_id, obs, history)

    try:
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt},
            ],
            temperature=0.05,
            max_tokens=120,
        )
        reply = response.choices[0].message.content.strip()

        # Strip markdown fences if the model wraps them
        if reply.startswith("```"):
            reply = reply.split("\n", 1)[-1].rsplit("```", 1)[0].strip()

        action_dict = json.loads(reply)
        return action_dict, None

    except Exception as e:
        error_msg = str(e)
        # Return fallback action with error
        return {"action_type": "notify_team", "target": "inference_service"}, error_msg

# ---------------------------------------------------------------------------
# Structured logging functions (hackathon format)
# ---------------------------------------------------------------------------
def log_start(task: str, env_name: str, model: str) -> None:
    """Log START with task, env, model."""
    print(f"[START] task={task} env={env_name} model={model}", flush=True)


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


def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
    """Log END with success, steps, final score, reward list."""
    success_val = str(success).lower()
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(
        f"[END] success={success_val} steps={steps} score={score:.3f} rewards={rewards_str}",
        flush=True,
    )


# ---------------------------------------------------------------------------
# Main inference loop
# ---------------------------------------------------------------------------
def run_inference():
    """Run all tasks and output structured logs."""
    env = ResilientAgentEnvironment()

    tasks = ["task1_latency_spike", "task2_prediction_drift", "task3_cascading_failure"]
    all_results = {}

    for task_id in tasks:
        task_short = task_id.split("_", 1)[1] if "_" in task_id else task_id
        env_name = "resilientagent-prod"
        
        # Reset environment
        obs = env.reset(task_id=task_id)
        
        # Log start
        log_start(task_short, env_name, MODEL_NAME)
        
        history: list[dict] = []
        step_rewards: List[float] = []
        max_steps = 10
        last_error = None

        # Run steps
        while len(history) < max_steps and not obs.done:
            step_num = len(history) + 1
            
            # Get LLM action
            action_dict, llm_error = get_llm_action(task_id, obs, history)
            
            action_type = action_dict.get("action_type", "check_metrics")
            target = action_dict.get("target", "inference_service")
            
            # Format action as string
            action_str = f"{action_type}('{target}')"
            
            # Execute action
            action = ResilientAgentAction(action_type=action_type, target=target)
            obs = env.step(action)
            
            # Record step
            step_rewards.append(obs.reward)
            last_error = llm_error
            
            log_step(step_num, action_str, obs.reward, obs.done, llm_error)
            
            history.append({
                "action_type": action_type,
                "target": target,
                "reward": obs.reward,
            })
            
            if obs.done:
                break
        
        # Grade task
        score = env.grade()
        success = env._model_healthy
        
        log_end(success, len(history), score, step_rewards)
        
        all_results[task_short] = {
            "score": round(score, 4),
            "steps": len(history),
            "resolved": success,
            "rewards": [round(r, 2) for r in step_rewards]
        }
    
    return all_results


if __name__ == "__main__":
    if not HF_TOKEN:
        print("WARNING: HF_TOKEN not set — API calls will likely fail.", file=sys.stderr)
    run_inference()