# Copyright 2024 Klira AI
#
# 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.

"""Evaluation runner — runs target functions against test datasets."""

from __future__ import annotations

import logging
import time
import uuid
from typing import Any, Callable

from opentelemetry import trace

from klira.sdk.decorators.core import OUTPUT_TRUNCATION_LIMIT
from klira.sdk.evals.dataset import load_dataset
from klira.sdk.evals.types import KliraEvalResult, KliraTestCase

logger = logging.getLogger(__name__)

_tracer = trace.get_tracer("klira.evals")


def evaluate(
    target: Callable[..., Any],
    *,
    data: str | list[KliraTestCase] | None = None,
    dataset_id: str | None = None,
    experiment_id: str | None = None,
    api_key: str | None = None,
    endpoint: str | None = None,
) -> KliraEvalResult:
    """Run an evaluation of a target function against a dataset.

    The SDK captures traces for each test case. The platform evaluates
    them server-side using an LLM judge.

    Args:
        target: The function to evaluate (typically a @workflow decorated function).
        data: Local dataset path (str) or pre-loaded test cases (list).
        dataset_id: Remote dataset ID (takes precedence over local data).
        experiment_id: Optional experiment identifier for grouping runs.
        api_key: API key for remote dataset fetching.
        endpoint: API endpoint for remote datasets.

    Returns:
        KliraEvalResult with test cases and timing info.

    Raises:
        ValueError: If no dataset source is provided.
    """
    run_id = experiment_id or f"eval_{uuid.uuid4().hex[:12]}"

    # Load test cases
    if isinstance(data, list):
        test_cases = data
    else:
        test_cases = load_dataset(
            path=data,
            dataset_id=dataset_id,
            api_key=api_key,
            endpoint=endpoint,
        )

    logger.info(
        "Starting evaluation run %s with %d test cases",
        run_id,
        len(test_cases),
    )

    start_time = time.monotonic()
    errors: list[str] = []

    for i, test_case in enumerate(test_cases):
        test_case_id = f"{run_id}_tc_{i}"
        try:
            _run_test_case(target, test_case, run_id, test_case_id, dataset_id)
        except Exception as e:
            error_msg = f"Test case {i} failed: {e}"
            logger.warning(error_msg)
            errors.append(error_msg)

    duration = time.monotonic() - start_time

    # Flush traces — in eval mode, failure should be reported
    _flush_traces()

    logger.info(
        "Evaluation run %s completed in %.2fs (%d/%d succeeded)",
        run_id,
        duration,
        len(test_cases) - len(errors),
        len(test_cases),
    )

    return KliraEvalResult(
        test_cases=test_cases,
        total_test_cases=len(test_cases),
        duration_seconds=duration,
        experiment_id=run_id,
        errors=errors,
    )


def _run_test_case(
    target: Callable[..., Any],
    test_case: KliraTestCase,
    run_id: str,
    test_case_id: str,
    dataset_id: str | None,
) -> Any:
    """Run a single test case with tracing."""
    attributes: dict[str, Any] = {
        "klira.entity_type": "eval_test_case",
        "klira.evals.evals_run": run_id,
        "klira.evals.test_case_id": test_case_id,
    }
    if dataset_id:
        attributes["klira.evals.dataset_id"] = dataset_id
    if test_case.expected_output:
        expected = test_case.expected_output
        if len(expected) > OUTPUT_TRUNCATION_LIMIT:
            expected = expected[:OUTPUT_TRUNCATION_LIMIT]
        attributes["klira.evals.expected_output"] = expected

    with _tracer.start_as_current_span("klira.evals.test_case", attributes=attributes):
        return target(test_case.input)


def _flush_traces() -> None:
    """Flush the tracer provider. Raises on failure in eval mode."""
    provider = trace.get_tracer_provider()
    if hasattr(provider, "force_flush"):
        try:
            provider.force_flush(timeout_millis=5000)
        except Exception as e:
            logger.error("Failed to flush eval traces: %s", e)
            raise RuntimeError(
                f"Failed to flush eval traces: {e}. Eval results may be incomplete."
            ) from e
