"""
inference.py
------------
Inference script — Smart Contract Audit RL Environment.

Implements agents for all three tasks using the OpenAI-compatible client.
Emits mandatory structured stdout in the OpenEnv format.

MANDATORY ENV VARS:
  API_BASE_URL   LLM API endpoint   (default: https://api.openai.com/v1)
  MODEL_NAME     Model identifier   (default: gpt-4o-mini)
  HF_TOKEN       API key / HF token

MANDATORY STDOUT FORMAT (per episode):
  [START] task=<id> env=smart-contract-audit model=<model>
  [STEP]  step=<n> action=<str> reward=<0.00> done=<true|false> error=<str|null>
  [END]   success=<true|false> steps=<n> score=<0.000> rewards=<r1,r2,...>

Usage:
  python inference.py

Output:
  Structured stdout per episode, plus baseline_scores.json summary.
"""

import asyncio
import json
import os
import sys
import time
from typing import Any, Dict, List, Optional

from openai import OpenAI

from server import Task1Environment, Task2Environment, Task3Environment
from env.schemas import Action, ActionType
from utils import T1_SYSTEM, T2_SYSTEM, T3_SYSTEM
from dotenv import load_dotenv

# ─────────────────────────────────────────────────────────────────────────────
# Configuration
# ─────────────────────────────────────────────────────────────────────────────

load_dotenv()  # Load from .env if available; otherwise rely on actual env vars
API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
MODEL_NAME   = os.getenv("MODEL_NAME",   "gpt-4o")
HF_TOKEN     = os.getenv("HF_TOKEN")

if not HF_TOKEN:
    print("[WARN] HF_TOKEN not set — API calls may fail.", file=sys.stderr)
    exit(1)

ENV_BENCHMARK          = "smart-contract-audit"
NUM_EPISODES           = 2      # keep low on free tier; raise for full eval
SEED_BASE              = 42

# Max LLM calls per episode (including the mandatory submit on last step).
# Budget: free tier handles ~5-6 calls per episode before rate-limiting.
MAX_STEPS_T1           = 5
MAX_STEPS_T2           = 4
MAX_STEPS_T3           = 4

# How many steps before the end we start injecting "submit now" pressure.
# E.g. PRESSURE_AT=2 means last 2 steps show a warning.
PRESSURE_AT            = 2

# Sliding-window size: how many recent (user, assistant) pairs to keep.
# system prompt + 2 exchanges = ~800 tokens max — safe for free tier.
HISTORY_WINDOW         = 2

# Truncate action results to this many chars before inserting into the prompt.
MAX_RESULT_CHARS       = 400

# A grader_score >= this threshold → success=true in [END] line
SUCCESS_SCORE_THRESHOLD = 0.5

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


# ─────────────────────────────────────────────────────────────────────────────
# Mandatory stdout helpers
# ─────────────────────────────────────────────────────────────────────────────

def log_start(task: str, env: str, model: str) -> None:
    """Emit the [START] line — one per episode."""
    print(f"[START] task={task} env={env} model={model}", flush=True)


def log_step(
    step: int,
    action: str,
    reward: float,
    done: bool,
    error: Optional[str] = None,
) -> None:
    """Emit a [STEP] line — one per env.step() call."""
    error_val = error if error else "null"
    print(
        f"[STEP] step={step} action={action} "
        f"reward={reward:.2f} done={str(done).lower()} error={error_val}",
        flush=True,
    )


def log_end(
    success: bool,
    steps: int,
    score: float,
    rewards: List[float],
) -> None:
    """Emit the [END] line — one per episode, always emitted."""
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(
        f"[END] success={str(success).lower()} steps={steps} "
        f"score={score:.3f} rewards={rewards_str}",
        flush=True,
    )


# ─────────────────────────────────────────────────────────────────────────────
# Shared utilities
# ─────────────────────────────────────────────────────────────────────────────

def _truncate(text: str, limit: int = MAX_RESULT_CHARS) -> str:
    """Truncate long action results to keep prompts small."""
    if len(text) <= limit:
        return text
    return text[:limit] + f"... [truncated, {len(text) - limit} chars omitted]"


def _sliding_messages(system: str, history: List[Dict[str, str]]) -> List[Dict[str, str]]:
    """
    Return system prompt + the last HISTORY_WINDOW (user, assistant) pairs.
    This keeps total tokens bounded regardless of episode length.
    """
    # history = [..., user, assistant, user, assistant, ...]
    # We want the last HISTORY_WINDOW complete pairs (2 messages each).
    keep = HISTORY_WINDOW * 2
    windowed = history[-keep:] if len(history) > keep else history
    return [{"role": "system", "content": system}] + windowed


def _call_llm(messages: List[Dict[str, str]], max_tokens: int = 150) -> tuple[str, Optional[str]]:
    """Call the LLM; return (raw_response, error_string_or_None)."""
    try:
        resp = client.chat.completions.create(
            model=MODEL_NAME, # type: ignore
            messages=messages, # type: ignore
            max_tokens=max_tokens,
            temperature=0.0,
        )
        return resp.choices[0].message.content.strip(), None    # type: ignore
    except Exception as e:
        return "", str(e)[:80]


def _parse_action(raw: str, fallback_at: ActionType,
                  fallback_params: Dict[str, Any]) -> tuple[ActionType, Dict[str, Any]]:
    """Parse LLM JSON response into (ActionType, params). Use fallback on failure."""
    try:
        parsed = json.loads(raw)
        return ActionType(parsed["action"]), parsed.get("params", {})
    except Exception:
        return fallback_at, fallback_params


def _pressure_suffix(steps_left: int) -> str:
    """Return an urgent suffix when the step budget is nearly exhausted."""
    if steps_left <= 0:
        return (
            "\n\n⚠️  FINAL STEP — you MUST submit your best answer RIGHT NOW.\n"
            "Do not browse further. Emit a submit action immediately."
        )
    if steps_left <= PRESSURE_AT:
        return (
            f"\n\n⚠️  Only {steps_left} step(s) remaining. "
            "You should submit your answer in the next step or two."
        )
    return ""


# ─────────────────────────────────────────────────────────────────────────────
# Task 1 — Targeted Vulnerability Detection
# ─────────────────────────────────────────────────────────────────────────────

def _t1_user(obs: Dict[str, Any], steps_left: int) -> str:
    result = _truncate(obs.get("last_action_result") or "Episode just started.")
    return (
        f"Contract: {obs['contract_name']} | {obs['contract_description'][:80]}\n"
        f"Step {obs['step_count']} | Reward: {obs['cumulative_reward']:.2f}\n"
        f"Last action: {obs['last_action'] or 'None'}\n"
        f"Result: {result}"
        + _pressure_suffix(steps_left)
    )


def _t1_force_submit(obs: Dict[str, Any], history: List[Dict[str, str]]) -> tuple[ActionType, Dict[str, Any]]:
    """
    Build a forced submission from what we already know.
    Strategy: ask the LLM one more time with an explicit 'submit NOW' mandate.
    If that fails, fall back to a heuristic.
    """
    mandate = (
        "Based on everything you have seen, submit your best answer NOW.\n"
        "Respond ONLY with this JSON (fill in the values):\n"
        '{"action":"submit","params":{"function_name":"<best_guess>","vulnerability_type":"<best_guess>"}}'
    )
    messages = _sliding_messages(T1_SYSTEM, history) + [{"role": "user", "content": mandate}]
    raw, _ = _call_llm(messages, max_tokens=80)
    at, params = _parse_action(raw, ActionType.SUBMIT,
                                {"function_name": "withdraw",
                                 "vulnerability_type": "reentrancy"})
    # Guarantee it's always a submit
    if at != ActionType.SUBMIT:
        at = ActionType.SUBMIT
        if "function_name" not in params:
            params["function_name"] = "withdraw"
        if "vulnerability_type" not in params:
            params["vulnerability_type"] = "reentrancy"
    return at, params


def _run_t1_episode(env: Task1Environment, seed: int, ep_num: int) -> Dict[str, Any]:
    r   = env.reset(seed=seed)
    obs = r.observation.model_dump()
    log_start(task="task1_vuln_detection", env=ENV_BENCHMARK, model=MODEL_NAME) # type: ignore

    history:      List[Dict[str, str]] = []
    step_rewards: List[float] = []
    grader_score  = 0.0
    steps_taken   = 0
    error_msg: Optional[str] = None

    try:
        for step in range(1, MAX_STEPS_T1 + 1):
            steps_left = MAX_STEPS_T1 - step
            is_last    = (step == MAX_STEPS_T1)

            if is_last:
                # Never waste the last step on browsing — force a submission
                at, params = _t1_force_submit(obs, history)
            else:
                user_msg = _t1_user(obs, steps_left)
                history.append({"role": "user", "content": user_msg})
                messages  = _sliding_messages(T1_SYSTEM, history)
                raw, error_msg = _call_llm(messages)
                history.append({"role": "assistant", "content": raw})
                at, params = _parse_action(raw, ActionType.LIST_FUNCTIONS, {})

            result   = env.step(Action(action_type=at, params=params))
            obs      = result.observation.model_dump()
            r_val    = result.reward.value
            done     = result.done

            step_rewards.append(r_val)
            steps_taken = step
            log_step(step=step, action=at.value, reward=r_val, done=done, error=error_msg)

            if done:
                v = r_val
                grader_score = 1.0 if v >= 4.9 else (0.5 if v >= 0.9 else 0.0)
                break

            if not is_last:
                time.sleep(0.5)

    finally:
        success = grader_score >= SUCCESS_SCORE_THRESHOLD
        log_end(success=success, steps=steps_taken, score=grader_score, rewards=step_rewards)

    return {"episode": ep_num, "seed": seed, "contract": obs["contract_name"],
            "grader_score": grader_score, "cumulative_reward": obs["cumulative_reward"]}


# ─────────────────────────────────────────────────────────────────────────────
# Task 2 — Property Discovery
# ─────────────────────────────────────────────────────────────────────────────


def _t2_user(obs: Dict[str, Any], steps_left: int) -> str:
    extra  = obs.get("extra", {})
    result = _truncate(obs.get("last_action_result") or "Episode just started.")
    return (
        f"Contract: {obs['contract_name']} | "
        f"Function: {extra.get('target_function','?')} ({extra.get('target_signature','')})\n"
        f"Step {obs['step_count']} | Reward: {obs['cumulative_reward']:.2f}\n"
        f"Last action: {obs['last_action'] or 'None'}\n"
        f"Result: {result}"
        + _pressure_suffix(steps_left)
    )


def _t2_force_submit(obs: Dict[str, Any], history: List[Dict[str, str]]) -> tuple[ActionType, Dict[str, Any]]:
    """Force a submit_property based on everything seen so far."""
    extra = obs.get("extra", {})
    fn    = extra.get("target_function", "this function")
    mandate = (
        f"You must now submit your best property for '{fn}'.\n"
        "Write 2-3 sentences covering: what state changes, what is transferred, revert conditions.\n"
        "Respond ONLY with:\n"
        '{"action":"submit_property","params":{"property":"<your property here>"}}'
    )
    messages = _sliding_messages(T2_SYSTEM, history) + [{"role": "user", "content": mandate}]
    raw, _ = _call_llm(messages, max_tokens=200)
    at, params = _parse_action(raw, ActionType.SUBMIT_PROPERTY, {})
    if at != ActionType.SUBMIT_PROPERTY or not params.get("property", "").strip():
        at     = ActionType.SUBMIT_PROPERTY
        params = {"property": (
            f"After a successful call to {fn}, the contract updates its internal state "
            f"according to the function's logic. Reverts if input conditions are not met."
        )}
    return at, params


def _run_t2_episode(env: Task2Environment, seed: int, ep_num: int) -> Dict[str, Any]:
    r   = env.reset(seed=seed)
    obs = r.observation.model_dump()
    fn  = obs["extra"].get("target_function", "?")
    log_start(task="task2_property_discovery", env=ENV_BENCHMARK, model=MODEL_NAME) # type: ignore

    history:      List[Dict[str, str]] = []
    step_rewards: List[float] = []
    grader_score  = 0.0
    steps_taken   = 0
    error_msg: Optional[str] = None

    try:
        for step in range(1, MAX_STEPS_T2 + 1):
            steps_left = MAX_STEPS_T2 - step
            is_last    = (step == MAX_STEPS_T2)

            if is_last:
                at, params = _t2_force_submit(obs, history)
            else:
                user_msg = _t2_user(obs, steps_left)
                history.append({"role": "user", "content": user_msg})
                messages  = _sliding_messages(T2_SYSTEM, history)
                raw, error_msg = _call_llm(messages, max_tokens=250)
                history.append({"role": "assistant", "content": raw})
                at, params = _parse_action(raw, ActionType.GET_FUNCTION_NATSPEC, {})

            result   = env.step(Action(action_type=at, params=params))
            obs      = result.observation.model_dump()
            r_val    = result.reward.value
            done     = result.done

            step_rewards.append(r_val)
            steps_taken = step
            log_step(step=step, action=at.value, reward=r_val, done=done, error=error_msg)

            if done:
                grader_score = round(r_val / 5.0, 3) if r_val > 0 else 0.0
                break

            if not is_last:
                time.sleep(0.5)

    finally:
        success = grader_score >= SUCCESS_SCORE_THRESHOLD
        log_end(success=success, steps=steps_taken, score=grader_score, rewards=step_rewards)

    return {"episode": ep_num, "seed": seed,
            "contract": obs["contract_name"], "function": fn,
            "grader_score": grader_score, "cumulative_reward": obs["cumulative_reward"]}


# ─────────────────────────────────────────────────────────────────────────────
# Task 3 — Rule Checker
# ─────────────────────────────────────────────────────────────────────────────

def _t3_user(obs: Dict[str, Any], steps_left: int) -> str:
    extra  = obs.get("extra", {})
    result = _truncate(obs.get("last_action_result") or "Episode just started.")
    return (
        f"Contract: {obs['contract_name']}\n"
        f"Property: {extra.get('property_english', '(none)')[:200]}\n"
        f"Step {obs['step_count']} | Reward: {obs['cumulative_reward']:.2f}\n"
        f"Last action: {obs['last_action'] or 'None'}\n"
        f"Result: {result}"
        + _pressure_suffix(steps_left)
    )


def _t3_force_submit(obs: Dict[str, Any], history: List[Dict[str, str]]) -> tuple[ActionType, Dict[str, Any]]:
    """Force a submit_function based on everything seen so far."""
    prop = obs.get("extra", {}).get("property_english", "")
    mandate = (
        f"Property: {prop[:200]}\n"
        "Based on everything you have seen, which function violates this property?\n"
        "Respond ONLY with:\n"
        '{"action":"submit_function","params":{"function_name":"<your_best_guess>"}}'
    )
    messages = _sliding_messages(T3_SYSTEM, history) + [{"role": "user", "content": mandate}]
    raw, _ = _call_llm(messages, max_tokens=80)
    at, params = _parse_action(raw, ActionType.SUBMIT_FUNCTION, {})
    if at != ActionType.SUBMIT_FUNCTION or not params.get("function_name", "").strip():
        # Heuristic fallback: scan property text for a function name mention
        fn_candidates = ["withdraw", "emergencyDrain", "buyTokens", "setPrice",
                         "bid", "finalize", "stake", "claimRewards"]
        prop_lower = prop.lower()
        chosen = next((fn for fn in fn_candidates if fn.lower() in prop_lower), "withdraw")
        at     = ActionType.SUBMIT_FUNCTION
        params = {"function_name": chosen}
    return at, params


def _run_t3_episode(env: Task3Environment, seed: int, ep_num: int) -> Dict[str, Any]:
    r   = env.reset(seed=seed)
    obs = r.observation.model_dump()
    log_start(task="task3_rule_checker", env=ENV_BENCHMARK, model=MODEL_NAME) # type: ignore

    history:      List[Dict[str, str]] = []
    step_rewards: List[float] = []
    grader_score  = 0.0
    steps_taken   = 0
    error_msg: Optional[str] = None

    try:
        for step in range(1, MAX_STEPS_T3 + 1):
            steps_left = MAX_STEPS_T3 - step
            is_last    = (step == MAX_STEPS_T3)

            if is_last:
                at, params = _t3_force_submit(obs, history)
            else:
                user_msg = _t3_user(obs, steps_left)
                history.append({"role": "user", "content": user_msg})
                messages  = _sliding_messages(T3_SYSTEM, history)
                raw, error_msg = _call_llm(messages)
                history.append({"role": "assistant", "content": raw})
                at, params = _parse_action(raw, ActionType.GET_PROPERTY_SPECIFICATION, {})

            result   = env.step(Action(action_type=at, params=params))
            obs      = result.observation.model_dump()
            r_val    = result.reward.value
            done     = result.done

            step_rewards.append(r_val)
            steps_taken = step
            log_step(step=step, action=at.value, reward=r_val, done=done, error=error_msg)

            if done:
                v = r_val
                grader_score = 1.0 if v >= 4.9 else (0.3 if v >= 1.0 else 0.0)
                break

            if not is_last:
                time.sleep(0.5)

    finally:
        success = grader_score >= SUCCESS_SCORE_THRESHOLD
        log_end(success=success, steps=steps_taken, score=grader_score, rewards=step_rewards)

    return {"episode": ep_num, "seed": seed, "contract": obs["contract_name"],
            "grader_score": grader_score, "cumulative_reward": obs["cumulative_reward"]}


# ─────────────────────────────────────────────────────────────────────────────
# Task runners
# ─────────────────────────────────────────────────────────────────────────────

def run_task1(n: int = NUM_EPISODES) -> Dict[str, Any]:
    print("\n" + "="*60, flush=True)
    print("TASK 1: Targeted Vulnerability Detection", flush=True)
    print("="*60, flush=True)
    env      = Task1Environment()
    episodes = [_run_t1_episode(env, SEED_BASE + i, i + 1) for i in range(n)]
    avg_s    = sum(e["grader_score"] for e in episodes) / n
    avg_r    = sum(e["cumulative_reward"] for e in episodes) / n
    print(f"\n  Avg grader score : {avg_s:.3f}", flush=True)
    print(f"  Avg cum reward   : {avg_r:.2f}", flush=True)
    return {"task_id": "task1_vuln_detection", "name": "Targeted Vulnerability Detection",
            "status": "active", "num_episodes": n, "episodes": episodes,
            "avg_grader_score": avg_s, "avg_cumulative_reward": avg_r}


def run_task2(n: int = NUM_EPISODES) -> Dict[str, Any]:
    print("\n" + "="*60, flush=True)
    print("TASK 2: Property Discovery", flush=True)
    print("="*60, flush=True)
    env      = Task2Environment()
    episodes = [_run_t2_episode(env, SEED_BASE + i, i + 1) for i in range(n)]
    avg_s    = sum(e["grader_score"] for e in episodes) / n
    avg_r    = sum(e["cumulative_reward"] for e in episodes) / n
    print(f"\n  Avg grader score : {avg_s:.3f}", flush=True)
    print(f"  Avg cum reward   : {avg_r:.2f}", flush=True)
    return {"task_id": "task2_property_discovery", "name": "Property Discovery",
            "status": "active", "num_episodes": n, "episodes": episodes,
            "avg_grader_score": avg_s, "avg_cumulative_reward": avg_r}


def run_task3(n: int = NUM_EPISODES) -> Dict[str, Any]:
    print("\n" + "="*60, flush=True)
    print("TASK 3: Rule Checker", flush=True)
    print("="*60, flush=True)
    env      = Task3Environment()
    episodes = [_run_t3_episode(env, SEED_BASE + i, i + 1) for i in range(n)]
    avg_s    = sum(e["grader_score"] for e in episodes) / n
    avg_r    = sum(e["cumulative_reward"] for e in episodes) / n
    print(f"\n  Avg grader score : {avg_s:.3f}", flush=True)
    print(f"  Avg cum reward   : {avg_r:.2f}", flush=True)
    return {"task_id": "task3_rule_checker", "name": "Rule Checker",
            "status": "active", "num_episodes": n, "episodes": episodes,
            "avg_grader_score": avg_s, "avg_cumulative_reward": avg_r}


# ─────────────────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────────────────

async def main() -> None:
    print("Smart Contract Audit RL Environment — Baseline Inference", flush=True)
    print(f"Model     : {MODEL_NAME}", flush=True)
    print(f"Base URL  : {API_BASE_URL}", flush=True)
    print(f"Episodes  : {NUM_EPISODES} per task  |  "
          f"Max steps: T1={MAX_STEPS_T1} T2={MAX_STEPS_T2} T3={MAX_STEPS_T3}", flush=True)
    print(f"Hist window: last {HISTORY_WINDOW} exchanges  |  "
          f"Result truncation: {MAX_RESULT_CHARS} chars", flush=True)

    t1 = run_task1(NUM_EPISODES)
    t2 = run_task2(NUM_EPISODES)
    t3 = run_task3(NUM_EPISODES)

    results = {
        "model": MODEL_NAME, "base_url": API_BASE_URL,
        "tasks": [t1, t2, t3],
    }
    overall = sum(t["avg_grader_score"] for t in results["tasks"]) / 3
    results["overall_avg_score"] = overall

    print("\n" + "="*60, flush=True)
    print("BASELINE SUMMARY", flush=True)
    print("="*60, flush=True)
    for t in results["tasks"]:
        print(f"  ✅ {t['name']:40s}: {t['avg_grader_score']:.3f}", flush=True)
    print(f"\n  Overall avg grader score: {overall:.3f}", flush=True)

    with open("baseline_scores.json", "w") as f:
        json.dump(results, f, indent=2)
    print("\n  Scores written to baseline_scores.json", flush=True)


if __name__ == "__main__":
    asyncio.run(main())