"""Baseline inference script for OpenEnv Customer Support environment.

Uses the OpenAI client to interact with an LLM and solve tasks.
Follows the STRICT stdout format required by OpenEnv validation.

Environment variables:
    API_BASE_URL  — Base URL for the OpenAI-compatible API
    MODEL_NAME    — Model identifier
    HF_TOKEN      — Hugging Face token (used as API key)
"""

from __future__ import annotations

import json
import os
import sys
import traceback

import requests
from openai import OpenAI

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

API_BASE_URL = os.environ.get("API_BASE_URL", "https://api.openai.com/v1")
MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-4o-mini")
HF_TOKEN = os.environ.get("HF_TOKEN", "")

ENV_BASE_URL = os.environ.get("ENV_BASE_URL", "http://localhost:7860")

# ── OpenAI client ─────────────────────────────────────────────────────

client = OpenAI(
    base_url=API_BASE_URL,
    api_key=HF_TOKEN or "sk-placeholder",
)

# ── Task configs ──────────────────────────────────────────────────────

TASK_CONFIGS = {
    "easy_classification": {
        "difficulty": "easy",
        "max_steps": 3,
        "ticket_count": 5,
    },
    "medium_response": {
        "difficulty": "medium",
        "max_steps": 5,
        "ticket_count": 3,
    },
    "hard_resolution": {
        "difficulty": "hard",
        "max_steps": 10,
        "ticket_count": 3,
    },
}


# ── Helper: call LLM ─────────────────────────────────────────────────

def call_llm(system_prompt: str, user_prompt: str) -> str:
    """Call the LLM and return the response text."""
    try:
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            temperature=0.0,
            max_tokens=1024,
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        return f"ERROR: {e}"


# ── Helper: env API calls ────────────────────────────────────────────

def env_reset(task_id: str, ticket_index: int = 0) -> dict:
    """Reset the environment."""
    resp = requests.post(
        f"{ENV_BASE_URL}/reset",
        json={"task_id": task_id, "ticket_index": ticket_index},
        timeout=30,
    )
    resp.raise_for_status()
    return resp.json()


def env_step(action: dict) -> dict:
    """Send an action to the environment."""
    resp = requests.post(
        f"{ENV_BASE_URL}/step",
        json=action,
        timeout=30,
    )
    resp.raise_for_status()
    return resp.json()


# ── Prompts ───────────────────────────────────────────────────────────

CLASSIFY_SYSTEM = """You are a customer support ticket classifier.
Classify the ticket into exactly ONE of these categories:
- delivery_delay
- damaged_product
- refund_request
- billing_error
- product_inquiry
- account_issue

Reply with ONLY the category label, nothing else."""

RESPOND_SYSTEM = """You are a professional customer support agent.
Write a helpful, empathetic response to the customer's issue.
Guidelines:
- Address the customer by their first name
- Acknowledge their frustration if they are upset
- Reference their order if applicable
- Outline clear next steps
- Be professional and warm
- Keep it concise (2-4 sentences)
Reply with ONLY the response text."""

ACTION_SYSTEM = """You are deciding the resolution action for a customer support ticket.
Based on the ticket details, choose ONE action and respond in JSON format:

Rules:
- For damaged products with high value (>$200) or angry customers: {"action": "escalate", "reason": "..."}
- For billing errors or refund requests: {"action": "refund", "amount": <order_value>}
- For delivery delays with critical priority: {"action": "escalate", "reason": "..."}
- For simple inquiries or low-severity issues: {"action": "resolve"}
- For any critical priority tickets: {"action": "escalate", "reason": "..."}

Reply with ONLY the JSON object."""


# ── Run a single task ─────────────────────────────────────────────────

def run_task(task_id: str, ticket_index: int = 0) -> dict:
    """Run one episode of a task and return results."""
    config = TASK_CONFIGS[task_id]
    difficulty = config["difficulty"]

    # Reset
    reset_data = env_reset(task_id, ticket_index)
    obs = reset_data["observation"]
    ticket = obs["ticket"]

    ticket_text = (
        f"Ticket ID: {ticket['ticket_id']}\n"
        f"Customer: {ticket['customer_name']}\n"
        f"Issue: {ticket['issue_text']}\n"
        f"Order Value: ${ticket['order_value']}\n"
        f"Priority: {ticket['priority']}\n"
        f"Sentiment: {ticket['sentiment']}"
    )

    rewards = []
    step_num = 0
    done = False
    score = 0.0
    last_error = None

    print(f"[START] task={task_id} env=customer-support model={MODEL_NAME}")

    try:
        # ── Step 1: Classify ──
        if not done:
            classification = call_llm(CLASSIFY_SYSTEM, ticket_text)
            classification = classification.strip().lower().replace(" ", "_").replace("-", "_")

            step_result = env_step({
                "action_type": "classify",
                "classification": classification,
            })
            step_num += 1
            reward_val = step_result["reward"]["total"]
            done = step_result["done"]
            score = step_result["observation"]["score"]
            rewards.append(reward_val)

            print(
                f"[STEP] step={step_num} action=classify "
                f"reward={reward_val:.2f} done={str(done).lower()} error=null"
            )

        # ── Step 2: Respond (medium + hard) ──
        if not done and difficulty in ("medium", "hard"):
            response = call_llm(
                RESPOND_SYSTEM,
                f"{ticket_text}\n\nClassification: {classification}",
            )

            step_result = env_step({
                "action_type": "respond",
                "response_text": response,
            })
            step_num += 1
            reward_val = step_result["reward"]["total"]
            done = step_result["done"]
            score = step_result["observation"]["score"]
            rewards.append(reward_val)

            print(
                f"[STEP] step={step_num} action=respond "
                f"reward={reward_val:.2f} done={str(done).lower()} error=null"
            )

        # ── Step 3: Take action (hard only) ──
        if not done and difficulty == "hard":
            action_response = call_llm(
                ACTION_SYSTEM,
                f"{ticket_text}\n\nClassification: {classification}",
            )

            # Parse the JSON action
            try:
                action_json = json.loads(action_response)
            except json.JSONDecodeError:
                # Fallback: try to extract action from text
                action_json = {"action": "escalate", "reason": "Unable to parse, escalating"}

            action_type = action_json.get("action", "resolve")
            step_action: dict = {"action_type": action_type}

            if action_type == "refund":
                step_action["refund_amount"] = action_json.get("amount", ticket["order_value"])
            elif action_type == "escalate":
                step_action["escalation_reason"] = action_json.get(
                    "reason", "Customer issue requires management attention"
                )

            step_result = env_step(step_action)
            step_num += 1
            reward_val = step_result["reward"]["total"]
            done = step_result["done"]
            score = step_result["observation"]["score"]
            rewards.append(reward_val)

            print(
                f"[STEP] step={step_num} action={action_type} "
                f"reward={reward_val:.2f} done={str(done).lower()} error=null"
            )

    except Exception as exc:
        last_error = str(exc)
        print(
            f"[STEP] step={step_num} action=error "
            f"reward=0.05 done=true error={last_error}"
        )
        rewards.append(0.05)

    # ── Summary ──
    success = done and score > 0.3
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    print(
        f"[END] success={str(success).lower()} steps={step_num} "
        f"score={score:.2f} rewards={rewards_str}"
    )

    return {
        "task_id": task_id,
        "ticket_index": ticket_index,
        "success": success,
        "steps": step_num,
        "score": score,
        "rewards": rewards,
    }


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

def main():
    """Run inference across all tasks and ticket pools."""
    all_results = []

    for task_id, config in TASK_CONFIGS.items():
        ticket_count = config["ticket_count"]
        for idx in range(ticket_count):
            try:
                result = run_task(task_id, ticket_index=idx)
                all_results.append(result)
            except Exception as exc:
                print(
                    f"[END] success=false steps=0 score=0.05 "
                    f"rewards=0.05"
                )
                traceback.print_exc(file=sys.stderr)

    # Summary
    print("\n--- Summary ---")
    for r in all_results:
        print(
            f"  {r['task_id']}[{r['ticket_index']}]: "
            f"score={r['score']:.2f} steps={r['steps']} "
            f"success={r['success']}"
        )


if __name__ == "__main__":
    main()