"""Adapters that connect Ember NON graphs with the upstream GEPA optimizer."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable, Iterable, Sequence

from ember.non import build_graph
from gepa.core.adapter import EvaluationBatch, GEPAAdapter

BenchmarkFn = Callable[[list[str], int, float], "EvaluationMetrics"]


@dataclass(slots=True)
class EvaluationMetrics:
    """Raw evaluation results returned by the benchmark callable."""

    accuracy: float
    latency_p50: float
    latency_p95: float
    latency_p99: float
    cost_per_call: float
    total_tokens_in: int
    total_tokens_out: int
    num_examples: int


@dataclass(slots=True)
class ScoreWeights:
    """Weights used to contract multi-objective metrics into a scalar score."""

    accuracy: float = 1.0
    latency: float = 1e-3
    cost: float = 1.0

    def score(self, metrics: EvaluationMetrics) -> float:
        return (
            self.accuracy * metrics.accuracy
            - self.latency * metrics.latency_p95
            - self.cost * metrics.cost_per_call
        )


@dataclass(slots=True)
class NonEvaluationRequest:
    """Single evaluation job executed by :class:`NonGEPAAdapter`."""

    dataset_id: str
    seed: int
    fidelity: float = 1.0


@dataclass(slots=True)
class NonTrajectory:
    """Trajectory information captured for reflective feedback."""

    spec_text: str
    request: NonEvaluationRequest
    metrics: EvaluationMetrics


def _canonicalise_spec(text: str) -> list[str]:
    lines = [line.strip() for line in text.splitlines()]
    return [line for line in lines if line]


def serialise_spec(spec: Sequence[str]) -> str:
    """Convert a NON spec sequence to canonical newline-separated text."""

    return "\n".join(item.strip() for item in spec)


class NonGEPAAdapter(GEPAAdapter[NonEvaluationRequest, NonTrajectory, dict[str, Any]]):
    """Adapter that allows GEPA to optimise NON graph specifications."""

    def __init__(
        self,
        *,
        benchmark_fn: BenchmarkFn,
        component_name: str = "graph_spec",
        score_weights: ScoreWeights | None = None,
    ) -> None:
        self.benchmark_fn = benchmark_fn
        self.component_name = component_name
        self.score_weights = score_weights or ScoreWeights()

    # ------------------------------------------------------------------
    # GEPAAdapter API
    # ------------------------------------------------------------------

    def evaluate(
        self,
        batch: list[NonEvaluationRequest],
        candidate: dict[str, str],
        capture_traces: bool = False,
    ) -> EvaluationBatch[NonTrajectory, dict[str, Any]]:
        if self.component_name not in candidate:
            raise ValueError(f"Candidate missing component '{self.component_name}'.")

        spec_text = candidate[self.component_name]
        spec = _canonicalise_spec(spec_text)

        # Validate syntax early to surface parser failures with actionable errors.
        build_graph(spec)

        outputs: list[dict[str, Any]] = []
        scores: list[float] = []
        trajectories: list[NonTrajectory] | None = [] if capture_traces else None

        for request in batch:
            metrics = self.benchmark_fn(spec, request.seed, request.fidelity)
            outputs.append(
                {
                    "dataset_id": request.dataset_id,
                    "seed": request.seed,
                    "fidelity": request.fidelity,
                    "metrics": metrics,
                }
            )
            scores.append(self.score_weights.score(metrics))

            if capture_traces:
                assert trajectories is not None
                trajectories.append(
                    NonTrajectory(
                        spec_text=spec_text,
                        request=request,
                        metrics=metrics,
                    )
                )

        return EvaluationBatch(outputs=outputs, scores=scores, trajectories=trajectories)

    def make_reflective_dataset(
        self,
        candidate: dict[str, str],
        eval_batch: EvaluationBatch[NonTrajectory, dict[str, Any]],
        components_to_update: list[str],
    ) -> dict[str, list[dict[str, Any]]]:
        if eval_batch.trajectories is None:
            raise ValueError("Trajectories were not captured; cannot build reflective dataset.")

        datasets: dict[str, list[dict[str, Any]]] = {}
        for component in components_to_update:
            datasets[component] = [
                _trajectory_to_feedback(component, trajectory, self.score_weights)
                for trajectory in eval_batch.trajectories
            ]
        return datasets

    propose_new_texts = None  # Filled by GEPA unless the caller overrides it.


def _trajectory_to_feedback(
    component: str,
    trajectory: NonTrajectory,
    weights: ScoreWeights,
) -> dict[str, Any]:
    metrics = trajectory.metrics
    score = weights.score(metrics)
    instructions = (
        "Optimise the NON graph specification to increase the aggregate score "
        "while keeping latency and cost low. Preserve syntactic validity so the "
        "graph can be parsed by ember.non.build_graph."
    )

    return {
        "Inputs": {
            "component": component,
            "dataset": trajectory.request.dataset_id,
            "seed": trajectory.request.seed,
            "fidelity": trajectory.request.fidelity,
            "current_spec": trajectory.spec_text,
        },
        "Generated Outputs": {
            "accuracy": metrics.accuracy,
            "latency_p95": metrics.latency_p95,
            "latency_p99": metrics.latency_p99,
            "cost_per_call": metrics.cost_per_call,
            "total_tokens_in": metrics.total_tokens_in,
            "total_tokens_out": metrics.total_tokens_out,
            "score": score,
        },
        "Feedback": instructions,
    }


def build_seed_candidate(
    spec: Iterable[str],
    *,
    component_name: str = "graph_spec",
) -> dict[str, str]:
    """Utility helper to assemble a GEPA candidate mapping."""

    return {component_name: serialise_spec(list(spec))}
