"""
This module provides a mixin class for running lm-eval harness evaluations
against SGLang servers
"""

import os
from contextlib import contextmanager
from pathlib import Path
from typing import Any

import numpy as np
import requests
import yaml

from sglang.test.test_utils import dump_metric


@contextmanager
def scoped_env_vars(new_env: dict[str, str] | None):
    """Context manager to temporarily set environment variables."""
    if not new_env:
        yield
        return

    old_values = {}
    new_keys = []

    try:
        for key, value in new_env.items():
            if key in os.environ:
                old_values[key] = os.environ[key]
            else:
                new_keys.append(key)
            os.environ[key] = str(value)
        yield
    finally:
        for key, value in old_values.items():
            os.environ[key] = value
        for key in new_keys:
            os.environ.pop(key, None)


class LMEvalMixin:
    """
    Mixin class for running lm-eval harness evaluations.
    """

    other_args: list[str] = []
    model_config_name: str = ""
    default_rtol: float = 0.08

    def test_lm_eval(self):
        """Run lm-eval evaluation and validate results."""
        # Flush cache before evaluation
        requests.get(self.base_url + "/flush_cache")

        eval_config = yaml.safe_load(
            Path(self.model_config_name).read_text(encoding="utf-8")
        )
        results = self.launch_lm_eval(eval_config)

        rtol = eval_config.get("rtol", self.default_rtol)

        success = True
        for task in eval_config["tasks"]:
            for metric in task["metrics"]:
                ground_truth = metric["value"]
                measured_value = results["results"][task["name"]][metric["name"]]
                print(
                    f"{task['name']} | {metric['name']}: "
                    f"ground_truth={ground_truth:.3f} | "
                    f"measured={measured_value:.3f} | rtol={rtol}"
                )
                dump_metric(
                    f"{task['name']}_{metric['name']}",
                    measured_value,
                    labels={
                        "model": eval_config.get("model_name", ""),
                        "eval": "lm-eval",
                        "task": task["name"],
                    },
                )
                success = success and np.isclose(
                    ground_truth, measured_value, rtol=rtol
                )

        self.assertTrue(success, f"lm-eval validation failed")

    def launch_lm_eval(self, eval_config: dict[str, Any]) -> dict:
        """
        Args:
            eval_config: Configuration dictionary with model and task settings
        """
        import lm_eval

        batch_size = eval_config.get("batch_size", "auto")
        backend = eval_config.get("backend", "local-completions")
        num_concurrent = eval_config.get("num_concurrent", 1)

        model_args = {
            "model": eval_config["model_name"],
            "base_url": self.base_url + "/v1/completions",
            "num_concurrent": num_concurrent,
        }

        env_vars = eval_config.get("env_vars", None)
        with scoped_env_vars(env_vars):
            results = lm_eval.simple_evaluate(
                model=backend,
                model_args=model_args,
                tasks=[task["name"] for task in eval_config["tasks"]],
                num_fewshot=eval_config.get("num_fewshot", 0),
                limit=eval_config.get("limit", None),
                apply_chat_template=eval_config.get("apply_chat_template", False),
                fewshot_as_multiturn=eval_config.get("fewshot_as_multiturn", False),
                gen_kwargs=eval_config.get("gen_kwargs"),
                batch_size=batch_size,
            )

        return results