import os
import json
import requests
from typing import Optional
from openai import OpenAI

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", "")
ENV_URL = os.getenv("ENV_URL", "https://dev9269-ai-support-ticket.hf.space")

TASKS = ["classify_ticket", "resolve_ticket", "triage_queue"]
MAX_STEPS = {"classify_ticket": 1, "resolve_ticket": 1, "triage_queue": 3}

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


def log_start(task, env, model):
    print(f"[START] task={task} env={env} model={model}", flush=True)


def log_step(step, action, reward, done, error: Optional[str] = None):
    print(f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()} error={error or 'null'}", flush=True)


def log_end(success, steps, score, rewards):
    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)


def call_llm(prompt: str) -> str:
    try:
        resp = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": "You are a customer support AI. Respond only with valid JSON."},
                {"role": "user", "content": prompt},
            ],
            temperature=0.3,
            max_tokens=300,
        )
        return resp.choices[0].message.content.strip()
    except Exception as e:
        return "{}"


def parse_action(text: str) -> dict:
    try:
        start = text.find("{")
        end = text.rfind("}") + 1
        return json.loads(text[start:end])
    except Exception:
        return {}


def run_task(task: str):
    log_start(task=task, env="ai-support-ticket", model=MODEL_NAME)

    obs = requests.post(f"{ENV_URL}/reset", json={"task": task}).json()
    rewards = []
    steps = 0
    score = 0.0
    success = False

    try:
        max_steps = MAX_STEPS.get(task, 1)
        for step in range(1, max_steps + 1):
            subject = obs.get("subject", "")
            description = obs.get("description", "")

            prompt = f"""Analyze this support ticket and respond with JSON:
Subject: {subject}
Description: {description}

Respond with:
{{
  "category": "<billing|technical|account|shipping|general>",
  "priority": "<low|medium|high|critical>",
  "resolution": "<your resolution text here>",
  "status": "closed"
}}"""

            raw = call_llm(prompt)
            action = parse_action(raw)

            result = requests.post(f"{ENV_URL}/step", json=action).json()
            reward = result.get("reward", 0.0)
            done = result.get("done", True)
            obs = result.get("observation", {})

            rewards.append(reward)
            steps = step
            score = result.get("info", {}).get("score", reward)

            log_step(step=step, action=json.dumps(action), reward=reward, done=done)

            if done:
                break

        success = score >= 0.5

    except Exception as e:
        log_step(step=steps + 1, action="{}", reward=0.0, done=True, error=str(e))

    finally:
        log_end(success=success, steps=steps, score=score, rewards=rewards)


if __name__ == "__main__":
    for task in TASKS:
        run_task(task)