from collections.abc import Callable
from typing import Any, Optional

from dspy.evaluate.evaluate import Evaluate
from dspy.teleprompt import bootstrap_trace_data
from dspy.teleprompt.bootstrap_trace import FailedPrediction
from dspy.teleprompt.gepa.gepa import GEPAFeedbackMetric
from dspy.teleprompt.gepa.gepa_utils import DSPyTrace
from models import (
    RPCLM,
    RPCModule,
    RPCRolloutResponse,
    RPCRolloutResult,
    RPCTraceStep,
)
from models import (
    Prediction as RPCPrediction,
)

import dspy
from dspy import Example, Prediction


def evaluate(
    module: dspy.Module,
    canidates: RPCModule | list[RPCModule],
    devset: list[dict[str, Any]],
    metric: Callable[[dspy.Example, dspy.Prediction, DSPyTrace | None], float]
    | GEPAFeedbackMetric,
    trace: bool = False,
    lm: RPCLM | None = None,
) -> RPCRolloutResponse:
    """Evaluate a DSPy program on a list of examples."""
    # NOTE: lm is currently unused as RPCLM conversion to dspy.LM is not yet implemented.
    if isinstance(canidates, list):
        msg = "Dspy harness does not support multiple candidates yet."
        raise NotImplementedError(msg)
    canidate = canidates

    module_copy = module.deepcopy()
    load_candidate(module_copy, canidate)

    dspy_devset = [
        Example(**example["inputs"], **(example.get("outputs") or {})).with_inputs(
            *list(example["inputs"].keys()),
        )
        for example in devset
    ]

    # Create name map for tracing
    name_map = {id(p): name for name, p in module_copy.named_predictors()}

    if trace:
        trace_data_list = bootstrap_trace_data(
            module_copy,
            dspy_devset,
            metric,
        )

        results = [
            RPCRolloutResult(
                example=dspy2rpc_example(trace_data["example"]),
                prediction=dspy2rpc_prediction(trace_data["prediction"]),
                score=trace_data.get("score"),
                trace=dspy2rpc_trace(trace_data["trace"], name_map),
            )
            for trace_data in trace_data_list
        ]
        return RPCRolloutResponse(results=results)

    evaluator = Evaluate(
        devset=dspy_devset,
        num_threads=None,
        display_progress=False,
    )

    eval_result = evaluator(module_copy, metric=metric)

    results = [
        RPCRolloutResult(
            example=dspy2rpc_example(ex),
            prediction=dspy2rpc_prediction(pred),
            score=score,
        )
        for ex, pred, score in eval_result.results
    ]
    return RPCRolloutResponse(results=results)


def load_candidate(module: dspy.Module, canidate: RPCModule) -> None:
    """Load candidate settings into the DSPy module."""
    for name, dspy_pred in module.named_predictors():
        if rpc_pred := canidate.named_predictors.get(name):
            new_sig = dspy_pred.signature.with_instructions(
                rpc_pred.signature.instructions,
            )
            dspy_pred.demos = [
                Example(**demo["inputs"], **(demo.get("outputs") or {})).with_inputs(
                    *demo["inputs"].keys(),
                )
                for demo in rpc_pred.demos
            ]
            # TODO: Convert RPCLM to dspy.LM if rpc_pred.lm is set
            # dspy_pred.lm = rpc_pred.lm
            dspy_pred.signature = new_sig
        else:
            msg = f"Predictor {name} not found in canidate {canidate.named_predictors}"
            raise ValueError(msg)


def dspy2rpc_example(obj: Example) -> dict[str, Any]:
    """Convert dspy.Example to RPCExample (dict)."""
    inputs = {k: obj[k] for k in obj._input_keys}
    outputs = {k: obj[k] for k in obj if k not in obj._input_keys}
    return {"inputs": inputs, "outputs": outputs}


def dspy2rpc_prediction(obj: Prediction | Any) -> RPCPrediction:
    """Convert dspy.Prediction to optiglot.models.Prediction."""
    if isinstance(obj, FailedPrediction):
        return RPCPrediction(
            output=None,
            error_message=obj.completion_text,
        )

    if isinstance(obj, Prediction):
        return RPCPrediction(output=dict(obj))

    try:
        # Some objects might be dict-like but not Prediction
        return RPCPrediction(output=dict(obj))
    except (TypeError, ValueError):
        return RPCPrediction(output=None, error_message=str(obj))


def dspy2rpc_trace(
    obj: list[tuple[Any, dict[str, Any], Prediction]],
    name_map: dict[int, str],
) -> list[RPCTraceStep]:
    """Convert DSPy trace to list of RPCTraceStep."""
    trace_steps = []
    for predictor, inputs, prediction in obj:
        predictor_name = name_map.get(id(predictor), "unknown")
        trace_steps.append(
            RPCTraceStep(
                predictor_name=predictor_name,
                inputs=inputs,
                prediction=dspy2rpc_prediction(prediction),
            ),
        )
    return trace_steps
