"""Temporal Activities: all nondeterministic work lives here.

Inference, wall-clock timing, and evaluation are side-effecting or
time-dependent, so they must not run inside the Workflow.
"""

import asyncio
import os
import time

from openai import OpenAI
from temporalio import activity

LLM_BASE_URL = os.getenv("LLM_BASE_URL", os.getenv("VLLM_BASE_URL", "http://127.0.0.1:8090/v1"))
LLM_API_KEY = os.getenv("LLM_API_KEY", os.getenv("VLLM_API_KEY", "EMPTY"))
MODEL_NAME = os.getenv("MODEL_NAME", "qwen38-27b")

# Durability / failure-injection knobs (sections 16 and 21 of the guide).
SLOW_ACTIVITY_SECONDS = float(os.getenv("SLOW_ACTIVITY_SECONDS", "0"))
FAIL_FIRST_ATTEMPT = os.getenv("FAIL_FIRST_ATTEMPT", "") == "1"

# Section 19: named backends so a configuration can select a model.
MODELS = {
    "default": {"base_url": LLM_BASE_URL, "api_key": LLM_API_KEY, "model": MODEL_NAME},
    "sglang-rtx6000": {
        "base_url": "http://127.0.0.1:8090/v1",
        "api_key": "VLLM_API_KEY",
        "model": "qwen38-27b",
    },
    "llamacpp-3090": {
        "base_url": "http://127.0.0.1:8080/v1",
        "api_key": "EMPTY",
        "model": "Qwen3.8-27B",
    },
}

_clients: dict = {}


def _get_client(base_url: str, api_key: str) -> OpenAI:
    key = (base_url, api_key)
    if key not in _clients:
        _clients[key] = OpenAI(base_url=base_url, api_key=api_key, timeout=300.0)
    return _clients[key]


@activity.defn
async def build_prompt(config: dict) -> dict:
    task = config["task"]
    context = config.get("context", "")

    system_prompt = config.get("system_prompt") or (
        "You are a precise engineering assistant. Answer concisely and accurately."
    )

    max_context_chars = config.get("max_context_chars")
    if max_context_chars is not None:
        context = context[:max_context_chars]

    user_prompt = f"TASK:\n{task}\n\nCONTEXT:\n{context}"

    return {"system_prompt": system_prompt, "user_prompt": user_prompt}


@activity.defn
async def call_llm(data: dict) -> dict:
    # Section 16: give the operator a window to kill the Worker mid-Activity.
    if SLOW_ACTIVITY_SECONDS > 0:
        deadline = time.monotonic() + SLOW_ACTIVITY_SECONDS
        while time.monotonic() < deadline:
            activity.heartbeat("waiting before inference")
            await asyncio.sleep(1)

    # Section 21: deterministic failure on the first attempt only.
    if FAIL_FIRST_ATTEMPT and activity.info().attempt == 1:
        raise RuntimeError("Injected test failure (attempt 1)")

    backend = MODELS.get(data.get("model") or "default", MODELS["default"])
    client = _get_client(backend["base_url"], backend["api_key"])
    model = backend["model"]

    extra_body = {}
    if data.get("enable_thinking") is not None:
        extra_body["chat_template_kwargs"] = {"enable_thinking": data["enable_thinking"]}

    started = time.perf_counter()
    response = await asyncio.to_thread(
        lambda: client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": data["system_prompt"]},
                {"role": "user", "content": data["user_prompt"]},
            ],
            temperature=data.get("temperature", 0.0),
            max_tokens=data.get("max_tokens", 256),
            extra_body=extra_body or None,
        )
    )
    elapsed = time.perf_counter() - started

    usage = response.usage
    message = response.choices[0].message
    completion = message.content or ""
    reasoning = getattr(message, "reasoning_content", None) or ""

    details = getattr(usage, "completion_tokens_details", None)
    reasoning_tokens = getattr(usage, "reasoning_tokens", None)
    if reasoning_tokens is None and details is not None:
        reasoning_tokens = getattr(details, "reasoning_tokens", None)

    return {
        "model": model,
        "backend": data.get("model") or "default",
        "answer": completion,
        "reasoning_chars": len(reasoning),
        "reasoning_tokens": reasoning_tokens,
        "finish_reason": response.choices[0].finish_reason,
        "latency_seconds": elapsed,
        "prompt_tokens": getattr(usage, "prompt_tokens", None),
        "completion_tokens": getattr(usage, "completion_tokens", None),
        "total_tokens": getattr(usage, "total_tokens", None),
        "attempt": activity.info().attempt,
    }


@activity.defn
async def evaluate_answer(data: dict) -> dict:
    answer = data["answer"] or ""
    expected_phrase = data["expected_phrase"]
    label_set = data.get("label_set")

    if label_set:
        # Strict single-label grading: the answer must name the expected label
        # and no competing label. A substring test alone would score "not
        # MEMORY, this is THERMAL" as a MEMORY pass.
        upper = answer.upper()
        mentioned = [label for label in label_set if label in upper]
        passed = mentioned == [expected_phrase]
        predicted = mentioned[0] if len(mentioned) == 1 else None
    else:
        passed = expected_phrase.lower() in answer.lower()
        predicted = None

    return {
        "passed": passed,
        "expected_phrase": expected_phrase,
        "predicted": predicted,
    }