"""Temporal Workflows: durable, deterministic orchestration only."""

from datetime import timedelta

from temporalio import workflow
from temporalio.common import RetryPolicy

with workflow.unsafe.imports_passed_through():
    from activities import build_prompt, call_llm, evaluate_answer


# Section 22. Kept deliberately small: an LLM call is not idempotent, so each
# retry is a real duplicate inference. Three attempts bounds the waste.
LLM_RETRY_POLICY = RetryPolicy(
    initial_interval=timedelta(seconds=1),
    maximum_interval=timedelta(seconds=10),
    maximum_attempts=3,
)

DETERMINISTIC_RETRY_POLICY = RetryPolicy(maximum_attempts=3)


async def _run_one(
    task: str,
    context: str,
    expected_phrase: str,
    config: dict,
    label_set: list | None = None,
) -> dict:
    prompt = await workflow.execute_activity(
        build_prompt,
        {
            "task": task,
            "context": context,
            "system_prompt": config.get("system_prompt"),
            "max_context_chars": config.get("max_context_chars"),
        },
        start_to_close_timeout=timedelta(seconds=30),
        retry_policy=DETERMINISTIC_RETRY_POLICY,
    )

    llm_result = await workflow.execute_activity(
        call_llm,
        {
            **prompt,
            "model": config.get("model"),
            "enable_thinking": config.get("enable_thinking"),
            "temperature": config.get("temperature", 0),
            "max_tokens": config.get("max_tokens", 256),
        },
        start_to_close_timeout=timedelta(minutes=5),
        heartbeat_timeout=timedelta(seconds=30),
        retry_policy=LLM_RETRY_POLICY,
    )

    evaluation = await workflow.execute_activity(
        evaluate_answer,
        {
            "answer": llm_result["answer"],
            "expected_phrase": expected_phrase,
            "label_set": label_set,
        },
        start_to_close_timeout=timedelta(seconds=30),
        retry_policy=DETERMINISTIC_RETRY_POLICY,
    )

    return {"name": config["name"], **llm_result, **evaluation}


@workflow.defn
class AgentOptimizationWorkflow:
    """Compares agent configurations on a single task (sections 7-14)."""

    @workflow.run
    async def run(self, experiment: dict) -> dict:
        results = []
        for config in experiment["configs"]:
            results.append(
                await _run_one(
                    experiment["task"],
                    experiment.get("context", ""),
                    experiment["expected_phrase"],
                    config,
                )
            )
        return {"task": experiment["task"], "results": results}


@workflow.defn
class BenchmarkWorkflow:
    """Runs every configuration across a suite of cases (sections 17-18)."""

    def __init__(self) -> None:
        self._done = 0
        self._total = 0

    @workflow.query
    def progress(self) -> dict:
        return {"completed": self._done, "total": self._total}

    @workflow.run
    async def run(self, suite: dict) -> dict:
        cases = suite["cases"]
        configs = suite["configs"]
        self._total = len(cases) * len(configs)

        rows = []
        for config in configs:
            for case in cases:
                result = await _run_one(
                    case["task"],
                    case.get("context", ""),
                    case["expected_phrase"],
                    config,
                    suite.get("label_set"),
                )
                result["case_id"] = case["id"]
                rows.append(result)
                self._done += 1

        return {"suite": suite.get("name", "suite"), "rows": rows}