"""Langfuse logging for checkpoint eval batches (Kaggle + local)."""
from __future__ import annotations

import os
from typing import Any


def init_langfuse():
    from langfuse import Langfuse, get_client

    Langfuse(
        public_key=os.environ["LANGFUSE_PUBLIC_KEY"],
        secret_key=os.environ["LANGFUSE_SECRET_KEY"],
        host=os.getenv("LANGFUSE_HOST")
        or os.getenv("LANGFUSE_BASE_URL")
        or "https://cloud.langfuse.com",
    )
    lf = get_client()
    if not lf.auth_check():
        raise RuntimeError("Langfuse auth_check failed")
    return lf


def log_checkpoint_eval(
    *,
    step: int,
    hf_revision: str,
    train_model_id: str,
    results: list[dict[str, Any]],
    wandb_run_id: str | None = None,
    trace_name: str = "kaggle-ckpt-eval",
) -> dict[str, Any]:
    """
    Write one Langfuse agent trace for a checkpoint eval suite.
    Each problem becomes a nested span; aggregate scores on the trace.
    MUST be called BEFORE pushing/saving the checkpoint artifact.
    """
    lf = init_langfuse()
    n = len(results)
    mean_pass = sum(float(r.get("pass_frac", 0.0)) for r in results) / n if n else 0.0
    mean_reward = sum(float(r.get("reward_total", 0.0)) for r in results) / n if n else 0.0

    with lf.start_as_current_observation(
        name=trace_name,
        as_type="evaluator",
        metadata={
            "step": step,
            "hf_revision": hf_revision,
            "train_model_id": train_model_id,
            "wandb_run_id": wandb_run_id,
            "phase": "ckpt_eval",
            "n_problems": n,
        },
        input={"step": step, "hf_revision": hf_revision, "n_problems": n},
    ) as root:
        for r in results:
            with lf.start_as_current_observation(
                name=f"task:{r.get('task_id')}",
                as_type="span",
                input={"prompt": r.get("prompt", "")[:2000]},
                metadata={"task_id": r.get("task_id")},
            ) as span:
                span.update(
                    output={
                        "completion": (r.get("completion") or "")[:2000],
                        "pass_frac": r.get("pass_frac"),
                        "reward_total": r.get("reward_total"),
                        "ok": r.get("pass_frac", 0) >= 1.0,
                    }
                )
        summary = {
            "mean_pass_frac": mean_pass,
            "mean_reward_total": mean_reward,
            "n_problems": n,
            "n_perfect": sum(1 for r in results if r.get("pass_frac", 0) >= 1.0),
            "step": step,
            "hf_revision": hf_revision,
        }
        root.update(output=summary)
        lf.score_current_trace(name="mean_pass_frac", value=mean_pass)
        lf.score_current_trace(name="mean_reward_total", value=mean_reward)
        lf.score_current_trace(name="n_perfect", value=float(summary["n_perfect"]))
        lf.score_current_trace(name="checkpoint_step", value=float(step))
        trace_id = lf.get_current_trace_id() or ""

    lf.flush()
    return {"langfuse_trace_id": trace_id, **summary}