# Copyright (c) Metis. All rights reserved.

"""Evals algorithm — single-pass agent evaluation with configurable parallelism.

Runs every dataset item through the agent exactly once, collects reward scores,
and reports aggregate metrics. No optimization, no reflection — just evaluate.

Usage via the Mantis Platform UI:
    1. Select algorithm: "Evals"
    2. Configure parallelism (how many rollouts run concurrently)
    3. Pick a dataset and agent
    4. Start the run

The algorithm enqueues rollouts in batches of ``parallelism``, waits for each
batch to complete, extracts scores, and logs per-batch and aggregate metrics.
"""

from __future__ import annotations

import asyncio
import logging
import uuid
from typing import Any, ClassVar, Dict, List, Optional

from mantisdk.algorithm.base import Algorithm
from mantisdk.algorithm.utils import with_llm_proxy, with_store
from mantisdk.emitter import find_final_reward
from mantisdk.types import Dataset, NamedResources, PromptTemplate

logger = logging.getLogger(__name__)


class Evals(Algorithm):
    """Single-pass evaluation algorithm.

    Runs each dataset item through the agent once, collects the reward score,
    and reports aggregate results. Parallelism controls how many rollouts are
    in-flight at the same time.
    """

    # Platform metadata — auto-discovered by registry
    name: ClassVar[str] = "evals"
    display_name: ClassVar[str] = "Evals"
    description: ClassVar[str] = "Single-pass agent evaluation"
    requirements: ClassVar[Dict[str, Any]] = {"gpu": False, "min_runners": 1}
    config_schema: ClassVar[Optional[Dict[str, Any]]] = {
        "type": "object",
        "properties": {
            "parallelism": {
                "type": "integer",
                "default": 5,
                "minimum": 1,
                "maximum": 50,
                "description": "Number of rollouts to run concurrently",
            },
            "rollout_timeout": {
                "type": "number",
                "default": 600.0,
                "minimum": 30.0,
                "description": "Timeout per rollout batch in seconds",
            },
            "val_fraction": {
                "type": "number",
                "default": 0.0,
                "minimum": 0.0,
                "maximum": 0.5,
                "description": "Fraction of dataset held out for validation (0 = use all items)",
            },
        },
    }

    def __init__(
        self,
        *,
        parallelism: int = 5,
        rollout_timeout: float = 600.0,
    ) -> None:
        super().__init__()
        self.parallelism = parallelism
        self.rollout_timeout = rollout_timeout

    @with_store
    @with_llm_proxy(required=False)
    async def run(
        self,
        llm_proxy,  # injected by @with_llm_proxy (may be None)
        store,  # injected by @with_store
        train_dataset: Optional[Dataset[Any]] = None,
        val_dataset: Optional[Dataset[Any]] = None,
    ) -> None:
        """Run single-pass evaluation on the dataset.

        Uses train_dataset as the evaluation set (the Platform UI sends
        the selected dataset as train_dataset). val_dataset is ignored.
        """
        dataset = train_dataset
        if dataset is None:
            raise ValueError("A dataset is required for evaluation.")

        store = self.get_store()
        assert store is not None

        initial_resources = self.get_initial_resources()
        if not initial_resources:
            raise ValueError("Initial resources must be set before running Evals.")

        # Build resources: initial resources + LLM proxy resource (if available)
        # The LLM proxy resource is required by agents using the @gepa.agent
        # or @rollout decorators that declare an `llm: LLM` parameter.
        resources: NamedResources = dict(initial_resources)
        if llm_proxy is not None:
            resources["llm"] = llm_proxy.as_resource()
            logger.info(f"Added LLM proxy resource: model={resources['llm'].model}")

        # Publish resources once
        resources_id = f"evals-{uuid.uuid4().hex[:8]}"
        await store.update_resources(resources_id, resources)
        logger.info(f"Published resources: {resources_id}")

        # Prepare dataset items
        items: List[Dict[str, Any]] = []
        for i, item in enumerate(dataset):
            if isinstance(item, dict):
                items.append(item)
            else:
                items.append({"task": item, "id": str(i)})

        total = len(items)
        logger.info(
            f"Starting evaluation: {total} items, parallelism={self.parallelism}, " f"timeout={self.rollout_timeout}s"
        )

        # Process in batches
        all_scores: List[Optional[float]] = []
        step = 0

        for batch_start in range(0, total, self.parallelism):
            batch = items[batch_start : batch_start + self.parallelism]
            batch_num = batch_start // self.parallelism + 1
            total_batches = (total + self.parallelism - 1) // self.parallelism

            logger.info(f"Batch {batch_num}/{total_batches}: " f"enqueuing {len(batch)} rollouts")

            # Enqueue rollouts for this batch
            rollout_ids: List[str] = []
            for item in batch:
                rollout = await store.enqueue_rollout(
                    input=item,
                    mode="test",
                    resources_id=resources_id,
                )
                rollout_ids.append(rollout.rollout_id)

            # Wait for all rollouts in the batch to complete
            try:
                completed = await store.wait_for_rollouts(
                    rollout_ids=rollout_ids,
                    timeout=self.rollout_timeout,
                )
            except Exception as e:
                logger.error(f"Batch {batch_num} failed: {e}")
                all_scores.extend([None] * len(batch))
                continue

            # Extract scores from completed rollouts
            batch_scores: List[Optional[float]] = []
            for rollout in completed:
                spans = await store.query_spans(rollout.rollout_id)
                reward = find_final_reward(spans)
                batch_scores.append(reward)

            all_scores.extend(batch_scores)

            # Log per-batch metrics
            valid_scores = [s for s in batch_scores if s is not None]
            batch_avg = sum(valid_scores) / len(valid_scores) if valid_scores else 0.0
            batch_completed = sum(1 for r in completed if r.status == "succeeded")

            step += 1
            if hasattr(store, "log"):
                store.log(
                    {
                        "batch_avg_score": batch_avg,
                        "batch_completed": batch_completed,
                        "batch_total": len(batch),
                    },
                    step=step,
                )

            logger.info(
                f"Batch {batch_num}/{total_batches} done: "
                f"avg_score={batch_avg:.3f}, "
                f"completed={batch_completed}/{len(batch)}"
            )

        # Log aggregate metrics
        valid_scores = [s for s in all_scores if s is not None]
        mean_score = sum(valid_scores) / len(valid_scores) if valid_scores else 0.0
        completion_rate = len(valid_scores) / total if total > 0 else 0.0

        step += 1
        if hasattr(store, "log"):
            store.log(
                {
                    "mean_score": mean_score,
                    "completion_rate": completion_rate,
                    "total_evaluated": len(valid_scores),
                    "total_items": total,
                },
                step=step,
            )

        logger.info(
            f"Evaluation complete: mean_score={mean_score:.3f}, "
            f"completion_rate={completion_rate:.1%}, "
            f"{len(valid_scores)}/{total} items scored"
        )
