"""
HARvestGym — Inference Script
==============================

Runs the RL agent (driven by an LLM via OpenAI client) through three tasks:
  1. har_classify_easy   — Template 1: list products in a category
  2. har_classify_medium — Template 3: add product to guest cart
  3. har_pipeline_hard   — Template 6: complete guest checkout

STDOUT FORMAT (strictly enforced by hackathon):
  [START] task=<task_name> env=<benchmark> model=<model_name>
  [STEP]  step=<n> action=<action_str> reward=<0.00> done=<true|false> error=<msg|null>
  [END]   success=<true|false> steps=<n> score=<score> rewards=<r1,r2,...,rn>

Usage:
  HF_TOKEN=hf_xxx uv run inference.py
  HF_TOKEN=hf_xxx MODEL_NAME=Qwen/Qwen2.5-72B-Instruct uv run inference.py
"""

import asyncio
import json
import os
import sys
import textwrap
from typing import Any, List, Optional

from openai import OpenAI

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

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")

if not HF_TOKEN:
    raise ValueError(
        "HF_TOKEN environment variable is required but not set. "
        "Export it with: export HF_TOKEN=hf_your_token_here"
    )

BENCHMARK = "harvgym"
MAX_STEPS = 20
TEMPERATURE = 0.7
MAX_TOKENS = 512
SUCCESS_SCORE_THRESHOLD = 0.5

# Task definitions for inference
TASKS = [
    {
        "task_name": "har_classify_easy",
        "template_id": 1,
        "description": "List products in the 'Gear' category",
        "app_base_url": "http://ec2-16-59-2-56.us-east-2.compute.amazonaws.com:7770/",
        "difficulty": "easy",
    },
    {
        "task_name": "har_classify_medium",
        "template_id": 3,
        "description": "Add 'Radiant Tee' to a guest cart",
        "app_base_url": "http://ec2-16-59-2-56.us-east-2.compute.amazonaws.com:7770/",
        "difficulty": "medium",
    },
    {
        "task_name": "har_pipeline_hard",
        "template_id": 6,
        "description": "Complete a guest checkout for 'Radiant Tee'",
        "app_base_url": "http://ec2-16-59-2-56.us-east-2.compute.amazonaws.com:7770/",
        "difficulty": "hard",
    },
]

# ---------------------------------------------------------------------------
# Logging helpers (hackathon format)
# ---------------------------------------------------------------------------

def log_start(task: str, env: str, model: str) -> None:
    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:
    error_val = error if error else "null"
    done_val = str(done).lower()
    # Sanitize action: no newlines
    action_clean = action.replace("\n", " ").replace("\r", "")[:200]
    print(
        f"[STEP] step={step} action={action_clean} reward={reward:.2f} done={done_val} error={error_val}",
        flush=True,
    )


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


# ---------------------------------------------------------------------------
# System prompt
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = textwrap.dedent("""
You are an API agent. Your goal is to complete a real-world task by calling the correct
sequence of HTTP API endpoints on a live web application.

You have exactly these tools available (output ONE tool call per turn as JSON):

1. browser_agent(task, url)
   → Discovers which API endpoints exist for this app. Call this FIRST and ONLY ONCE.
   → Returns: list of {method, path} endpoint names (no schemas)

2. search_endpoints(query)
   → Semantic search for endpoint schemas. Use after browser_agent to get full details.
   → Example: search_endpoints("create guest cart") returns method, path, auth, params

3. curl_exec(command)
   → Execute an HTTP call. Returns {status_code, headers, body}.
   → Use full curl syntax: curl -X POST 'URL' -H 'Content-Type: application/json' -d '{...}'
   → Session cookies are auto-injected; you do NOT need to set Cookie headers manually.

4. search_episode_data(query)
   → Search all prior API responses in this episode for a specific value.
   → Use when a response list was truncated and you need a specific item.

5. done(result?)
   → Call when the task is complete.

RULES:
- Output ONLY a single JSON object with keys "tool" and "args". Nothing else.
- Call browser_agent exactly once at step 1.
- Read values from prior responses (cart_id, sku, tokens) from the history.
- For Magento Shopping API (port 7770/7780): use Content-Type: application/json
- For Forum Postmill (port 9999): use Content-Type: application/x-www-form-urlencoded for login/post
- For Wikipedia (port 8888): GET requests only

EXAMPLE output format:
{"tool": "curl_exec", "args": {"command": "curl -X POST 'http://ec2-.../rest/V1/guest-carts' -H 'Content-Type: application/json'"}}
""").strip()


# ---------------------------------------------------------------------------
# LLM agent loop
# ---------------------------------------------------------------------------

def build_user_prompt(task_desc: str, app_base_url: str, step: int,
                       last_result: Any, history: List[dict],
                       session_state: dict) -> str:
    """Build the user prompt for each step."""
    history_str = ""
    if history:
        recent = history[-6:]  # Last 6 steps to stay within context
        lines = []
        for h in recent:
            result_str = json.dumps(h.get("result", ""))[:500]
            lines.append(f"  Step {h['step']}: {h['tool']}({h.get('args', {})}) → {result_str}")
        history_str = "\n".join(lines)

    session_str = json.dumps(session_state, indent=2)[:300] if session_state else "{}"

    last_result_str = json.dumps(last_result)[:800] if last_result is not None else "null"

    return textwrap.dedent(f"""
    TASK: {task_desc}
    APP URL: {app_base_url}
    STEP: {step}/{MAX_STEPS}

    SESSION STATE (auto-managed cookies/tokens):
    {session_str}

    LAST TOOL RESULT:
    {last_result_str}

    RECENT HISTORY:
    {history_str if history_str else "  (none yet)"}

    What is your next tool call? Output ONLY the JSON object.
    """).strip()


def get_model_action(client: OpenAI, task_desc: str, app_base_url: str,
                     step: int, last_result: Any, history: List[dict],
                     session_state: dict) -> dict:
    """Ask the LLM for the next action. Returns parsed tool call dict."""
    user_prompt = build_user_prompt(task_desc, app_base_url, step,
                                    last_result, history, session_state)
    try:
        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            temperature=TEMPERATURE,
            max_tokens=MAX_TOKENS,
            stream=False,
        )
        text = (completion.choices[0].message.content or "").strip()

        # Parse JSON from response
        # Handle markdown code blocks
        if "```json" in text:
            text = text.split("```json")[1].split("```")[0].strip()
        elif "```" in text:
            text = text.split("```")[1].split("```")[0].strip()

        # Find first { ... } block
        start = text.find("{")
        end = text.rfind("}") + 1
        if start >= 0 and end > start:
            text = text[start:end]

        parsed = json.loads(text)
        if "tool" in parsed:
            return parsed

        # LLM returned something else — default to done
        return {"tool": "done", "args": {"result": "Model returned non-tool response"}}

    except json.JSONDecodeError:
        # Couldn't parse JSON — try to extract tool name at minimum
        if "browser_agent" in text:
            return {"tool": "browser_agent", "args": {"task": task_desc, "url": app_base_url}}
        elif "done" in text.lower():
            return {"tool": "done", "args": {}}
        else:
            return {"tool": "done", "args": {"result": f"Parse error: {text[:100]}"}}
    except Exception as exc:
        print(f"[DEBUG] LLM call failed: {exc}", flush=True)
        # Default to browser_agent on first step, done otherwise
        if step == 1:
            return {"tool": "browser_agent", "args": {"task": task_desc, "url": app_base_url}}
        return {"tool": "done", "args": {"result": f"LLM error: {exc}"}}


# ---------------------------------------------------------------------------
# Single task episode runner
# ---------------------------------------------------------------------------

async def run_episode(task_config: dict, client: OpenAI) -> dict:
    """
    Run a single episode for one task.

    Returns: {"task_name", "success", "steps", "score", "rewards"}
    """
    from server.models import HARvestGymEnvironment, HarvestGymAction

    task_name = task_config["task_name"]
    template_id = task_config["template_id"]
    task_description = task_config["description"]
    app_base_url = task_config["app_base_url"]

    # Configure environment for this task
    os.environ["HARVGYM_TASK"] = str(template_id)

    env = HARvestGymEnvironment()

    rewards: List[float] = []
    steps_taken = 0
    score = 0.0
    success = False
    last_result = None
    history: List[dict] = []
    session_state: dict = {}

    log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)

    try:
        # Reset
        obs = env.reset()
        task_desc = obs.task or task_description
        base_url = obs.app_base_url or app_base_url

        for step in range(1, MAX_STEPS + 1):
            if getattr(obs, "done", False):
                break

            # Get action from LLM
            action_dict = get_model_action(
                client=client,
                task_desc=task_desc,
                app_base_url=base_url,
                step=step,
                last_result=last_result,
                history=history,
                session_state=session_state,
            )

            tool = action_dict.get("tool", "done")
            args = action_dict.get("args", {})

            action_str = f"{tool}({json.dumps(args)[:150]})"
            error_str = None

            try:
                action = HarvestGymAction(tool=tool, args=args)
                obs = env.step(action)

                reward = float(obs.reward or 0.0)
                done = bool(obs.done)
                last_result = obs.last_tool_result
                session_state = dict(obs.session_state or {})

                # Update history
                history.append({
                    "step": step,
                    "tool": tool,
                    "args": args,
                    "result": last_result,
                })

            except Exception as exc:
                reward = -0.1
                done = False
                error_str = str(exc)[:200]

            rewards.append(reward)
            steps_taken = step
            log_step(step=step, action=action_str, reward=reward, done=done, error=error_str)

            if done:
                break

        # Compute episode score from cumulative rewards
        # Normalize: terminal reward dominates; clamp to [0, 1]
        total_reward = sum(rewards)
        # Map reward to [0, 1]: reward range is roughly [-1.5, +7.5] per design
        score = (total_reward + 1.5) / 9.0
        score = max(0.0, min(1.0, score))
        success = score >= SUCCESS_SCORE_THRESHOLD

    except Exception as exc:
        error_str = str(exc)[:200]
        print(f"[DEBUG] Episode error: {error_str}", flush=True)
    finally:
        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)

    return {
        "task_name": task_name,
        "success": success,
        "steps": steps_taken,
        "score": score,
        "rewards": rewards,
    }


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

async def main() -> None:
    client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)

    results = []
    for task_config in TASKS:
        result = await run_episode(task_config, client)
        results.append(result)

    # Summary
    print("\n[SUMMARY]", flush=True)
    for r in results:
        status = "PASS" if r["success"] else "FAIL"
        print(
            f"  [{status}] {r['task_name']} — score={r['score']:.2f} steps={r['steps']}",
            flush=True,
        )

    overall_score = sum(r["score"] for r in results) / len(results) if results else 0.0
    print(f"\n  overall_score={overall_score:.2f}", flush=True)


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