import os
from openai import OpenAI

# ============================================
# SETUP — Read environment variables
# ============================================

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

if HF_TOKEN is None:
    raise ValueError("HF_TOKEN environment variable is required")

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

# ============================================
# TASK EMAIL DATASETS
# ============================================

TASKS = {
    "easy_sorting": [
        {"subject": "You won $1,000,000!", "body": "Click here to claim your prize now.", "sender": "prize@randomsite.xyz", "label": "spam"},
        {"subject": "FREE iPhone giveaway", "body": "You have been selected. Claim before midnight!", "sender": "free@giveaway123.com", "label": "spam"},
        {"subject": "Congratulations! You're a winner", "body": "Send your details to collect your reward.", "sender": "win@prizes.net", "label": "spam"},
        {"subject": "Meeting at 3pm today", "body": "Hi, reminder about our team sync at 3pm.", "sender": "manager@company.com", "label": "important"},
        {"subject": "Your invoice is ready", "body": "Please find your monthly invoice attached.", "sender": "billing@service.com", "label": "important"},
    ],
    "medium_sorting": [
        {"subject": "50% off this weekend only!", "body": "Flash sale on all items. Use code SAVE50.", "sender": "deals@amazon.com", "label": "promotion"},
        {"subject": "New arrivals just for you", "body": "Check out our latest summer collection.", "sender": "news@shop.com", "label": "promotion"},
        {"subject": "Exclusive member offer inside", "body": "As a valued member, enjoy 20% off.", "sender": "offers@store.com", "label": "promotion"},
        {"subject": "Action required: password expiry", "body": "Your password expires in 3 days. Reset it now.", "sender": "security@company.com", "label": "important"},
        {"subject": "RE: Project update", "body": "Thanks for the update. Let's connect tomorrow.", "sender": "colleague@work.com", "label": "important"},
        {"subject": "Urgent: verify your account", "body": "Your account will be suspended. Click to verify.", "sender": "alert@bank-secure.xyz", "label": "spam"},
    ],
    "hard_sorting": [
        {"subject": "Your account statement", "body": "Your monthly statement from XYZ Bank is ready.", "sender": "statements@xyzbank.com", "label": "important"},
        {"subject": "Limited time: upgrade your plan", "body": "Switch to premium and save 30% this month only.", "sender": "offers@service.com", "label": "promotion"},
        {"subject": "Security alert", "body": "A new login was detected from an unknown device.", "sender": "security@google.com", "label": "important"},
        {"subject": "You have unclaimed rewards", "body": "Collect your loyalty points before they expire.", "sender": "rewards@airline.com", "label": "promotion"},
        {"subject": "Final notice: payment overdue", "body": "Send $500 to avoid service interruption.", "sender": "billing@suspicious.xyz", "label": "spam"},
        {"subject": "Team offsite next Friday", "body": "Please confirm your attendance for the offsite.", "sender": "hr@company.com", "label": "important"},
        {"subject": "Claim your free trial", "body": "Start your 30-day free trial — no credit card needed.", "sender": "trial@software.com", "label": "promotion"},
        {"subject": "You've been pre-approved!", "body": "You qualify for a $50,000 loan. Apply now.", "sender": "loans@quickcash.xyz", "label": "spam"},
    ],
}

# ============================================
# LLM CLASSIFIER
# ============================================

def classify_email(email: dict) -> str:
    prompt = f"""Classify this email into exactly one category: spam, important, or promotion.

Subject: {email['subject']}
From: {email['sender']}
Body: {email['body']}

Reply with ONLY one word: spam, important, or promotion"""

    try:
        response = client.chat.completions.create(
            model=MODEL_NAME,
            messages=[
                {"role": "system", "content": "You are an email classifier. Reply with only one word: spam, important, or promotion."},
                {"role": "user",   "content": prompt}
            ],
            max_tokens=10,
            temperature=0.0
        )
        answer = response.choices[0].message.content.strip().lower()
        if "spam"      in answer: return "spam"
        if "promotion" in answer: return "promotion"
        if "important" in answer: return "important"
        return "spam"
    except Exception:
        return fallback_classify(email)


def fallback_classify(email: dict) -> str:
    text     = (email["subject"] + " " + email["body"]).lower()
    spam_kw  = ["won", "free", "prize", "claim", "urgent", "congratulations",
                "selected", "suspended", "verify", "overdue", "pre-approved"]
    promo_kw = ["off", "sale", "deal", "discount", "offer", "save",
                "upgrade", "trial", "rewards", "loyalty"]
    spam_score  = sum(1 for k in spam_kw  if k in text)
    promo_score = sum(1 for k in promo_kw if k in text)
    if spam_score >= 2:  return "spam"
    if promo_score >= 1: return "promotion"
    return "important"

# ============================================
# RUN ONE TASK EPISODE
# ============================================

def run_task(task_id: str, emails: list):
    rewards = []
    success = True

    # [START] — one per task
    print(f"[START] task={task_id} env=email-sorting-openenv model={MODEL_NAME}", flush=True)

    try:
        for i, email in enumerate(emails, 1):
            action = classify_email(email)
            correct = (action == email["label"])
            reward  = 1.0 if correct else -0.5
            done    = (i == len(emails))

            rewards.append(reward)
            # [STEP] — one per email
            print(f"[STEP] step={i} action={action} reward={reward:.2f} done={'true' if done else 'false'} error=null", flush=True)

    except Exception as e:
        success = False

    correct_count = sum(1 for r in rewards if r > 0)
    score = round(correct_count / len(emails), 2) if emails else 0.0
    # Clamp strictly between 0 and 1
    score = round(min(0.99, max(0.01, score)), 2)

    rewards_str = ",".join(f"{r:.2f}" for r in rewards)

    # [END] — includes score= which validator reads for grader check
    print(f"[END] task={task_id} success={'true' if success else 'false'} steps={len(rewards)} score={score} rewards={rewards_str}", flush=True)

    return score

# ============================================
# ENTRY POINT — runs all 3 tasks
# ============================================

if __name__ == "__main__":
    for task_id, emails in TASKS.items():
        run_task(task_id, emails)