# Copyright 2025 Horizon RL Contributors

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Evaluator for running agentic benchmarks with `strands-env` environments."""

from __future__ import annotations

import asyncio
import json
import logging
from collections import defaultdict
from collections.abc import Awaitable, Callable, Iterable
from functools import partial
from pathlib import Path

from pydantic import BaseModel
from tqdm import tqdm
from tqdm.contrib.logging import logging_redirect_tqdm

from strands_env.core import Action, Environment, StepResult

from .metrics import MetricFn, pass_at_k_metric

logger = logging.getLogger(__name__)

#: Type alias for environment factory function (async).
AsyncEnvFactory = Callable[[Action], Awaitable[Environment]]


class EvalSample(BaseModel):
    """Evaluation sample result."""

    action: Action
    """The action (task) that was evaluated."""

    step_result: StepResult
    """The result of the step (observation, reward, termination reason)."""


class Evaluator:
    """Evaluator for running concurrent environment evaluations."""

    benchmark_name: str = ""
    """Benchmark identifier. Override in subclasses."""

    def __init__(
        self,
        env_factory: AsyncEnvFactory,
        *,
        max_concurrency: int = 10,
        n_samples_per_prompt: int = 1,
        output_path: Path | str = Path.cwd() / "results.jsonl",
        save_interval: int = 10,
        keep_tokens: bool = False,
        metric_fns: list[MetricFn] = [],
    ):
        """Initialize the evaluator.

        Args:
            env_factory: Async factory function that creates a fresh Environment per sample.
            max_concurrency: Maximum concurrent evaluate_sample() calls.
            n_samples_per_prompt: Number of samples per prompt (for pass@k, set to max(k_values)).
            output_path: Path to JSONL file for saving results. Enables resume.
            save_interval: Flush results to disk every N completed samples.
            keep_tokens: Keep token-level observation in results (only valid for `SGLangModel` backends).
            metric_fns: Additional metric functions. `pass@k` is always included.
        """
        self.env_factory: AsyncEnvFactory = env_factory
        self.max_concurrency = max_concurrency
        self.n_samples_per_prompt = n_samples_per_prompt
        self.output_path = Path(output_path)
        self.save_interval = save_interval
        self.keep_tokens = keep_tokens

        # Always include pass@k, then any additional metrics
        self.metric_fns: list[MetricFn] = [
            partial(pass_at_k_metric, k_values=list(range(1, n_samples_per_prompt + 1)), reward_threshold=1.0)
        ]
        self.metric_fns += metric_fns

        # Runtime state
        self.results: dict[str, list[EvalSample]] = defaultdict(list)
        self.completed_ids: set[str] = set()

    def load_dataset(self) -> Iterable[Action]:
        """Load dataset. Override in subclasses."""
        raise NotImplementedError("Subclasses must implement load_dataset()")

    def load_results(self) -> None:
        """Load completed samples from checkpoint file."""
        if not self.output_path.exists():
            return

        self.results = defaultdict(list)
        self.completed_ids = set()

        with open(self.output_path, encoding="utf-8") as f:
            for line in f:
                data = json.loads(line)
                prompt_id = data.pop("prompt_id")
                sample = EvalSample.model_validate(data)
                self.results[prompt_id].append(sample)
                self.completed_ids.add(sample.action.task_context.id)

        total = sum(len(s) for s in self.results.values())
        logger.info(f"Resumed {total} samples from {self.output_path}")

    def save_results(self) -> None:
        """Save all samples to checkpoint file."""
        self.output_path.parent.mkdir(parents=True, exist_ok=True)
        with open(self.output_path, "w", encoding="utf-8") as f:
            for prompt_id, samples in self.results.items():
                for sample in samples:
                    data = sample.model_dump()
                    data["prompt_id"] = prompt_id
                    f.write(json.dumps(data, ensure_ascii=False) + "\n")

    async def evaluate_sample(self, action: Action) -> EvalSample:
        """Evaluate a single sample."""
        env = await self.env_factory(action)
        await env.reset()
        step_result = await env.step(action)
        if not self.keep_tokens:
            step_result.observation.tokens = None
        await env.cleanup()
        # Runtime logging for debugging
        reward_str = f"{step_result.reward.reward:.2f}" if step_result.reward else "N/A"
        reward_info = step_result.reward.info if step_result.reward else {}
        logger.info(
            f"[{action.task_context.id}]: "
            f"reward={reward_str} | "
            f"label={action.task_context.ground_truth} | "
            f"reward_info={reward_info} | "
            f"metrics={step_result.observation.metrics}"
        )
        return EvalSample(action=action, step_result=step_result)

    async def run(self, actions: Iterable[Action]) -> dict[str, list[EvalSample]]:
        """Run evaluation on actions with n_samples_per_prompt each.

        Args:
            actions: Actions to evaluate.

        Returns:
            Dict mapping prompt_id to list of EvalSample results.
        """
        self.load_results()

        # Expand actions to (prompt_id, sample_id, action) tuples
        to_process: list[tuple[str, str, Action]] = []
        for action in actions:
            prompt_id = action.task_context.id
            for i in range(self.n_samples_per_prompt):
                sample_id = f"{prompt_id}_{i}"
                if sample_id not in self.completed_ids:
                    expanded = action.model_copy(deep=True)
                    expanded.task_context.id = sample_id
                    to_process.append((prompt_id, sample_id, expanded))

        semaphore = asyncio.Semaphore(self.max_concurrency)
        save_counter = 0
        total = len(to_process)

        async def process(prompt_id: str, sample_id: str, action: Action, pbar: tqdm) -> None:
            nonlocal save_counter
            async with semaphore:
                sample = await self.evaluate_sample(action)
                self.results[prompt_id].append(sample)
                self.completed_ids.add(sample_id)
                pbar.update(1)
                save_counter += 1
                if save_counter >= self.save_interval:
                    self.save_results()
                    save_counter = 0

        with logging_redirect_tqdm():
            with tqdm(total=total, desc=f"Evaluating {self.benchmark_name}", unit="sample", dynamic_ncols=True) as pbar:
                await asyncio.gather(*[process(pid, sid, a, pbar) for pid, sid, a in to_process])
        self.save_results()
        return dict(self.results)

    def compute_metrics(self, results: dict[str, list[EvalSample]], log: bool = True) -> dict[str, float]:
        """Compute all metrics on results.

        Args:
            results: Dict mapping prompt_id to sample results.
            log: Whether to log the metrics summary.

        Returns:
            Dict mapping metric names to values.
        """
        metrics = {}
        for fn in self.metric_fns:
            metrics.update(fn(results))

        if log and metrics:
            n_prompts = len(results)
            n_samples = sum(len(s) for s in results.values())
            name = self.benchmark_name or "Evaluation"

            # Build formatted output
            lines = [f"{'─' * 40}", f"  {name} Results", f"{'─' * 40}"]
            lines.append(f"  Prompts: {n_prompts}  Samples (n={self.n_samples_per_prompt}): {n_samples}")
            lines.append("")
            for metric, value in sorted(metrics.items()):
                lines.append(f"  {metric:<12} {value:>6.1%}")
            lines.append(f"{'─' * 40}")
            logger.info("\n" + "\n".join(lines))

        return metrics
