# replay_runner.py — re-run ONE instrumented entrypoint on a given input so it emits a fresh trace.
#
# Generated by the agent-observability-replay-trace skill. The skill invokes this (it is NOT a server):
#   DD_TAGS=replay_run_id:<marker> python replay_runner.py --entrypoint <id> --input-file <path.json>
# It runs the entrypoint locally and lets it emit its normal trace — structurally IDENTICAL to a regular
# run, no wrapper span. Correlation is via the DD_TAGS=replay_run_id:<marker> the caller sets in this
# process's environment: ddtrace stamps it on every span of the run (the same channel that carries
# git_commit_sha / env), so the caller can find the new trace by that tag without altering trace shape.
#
# Fill the {{PLACEHOLDERS}} and the ENTRYPOINTS table (one entry per execution type). Keep this file and
# extend the table as new entrypoints appear.
#
# Requires in .env (loaded before ddtrace imports): DD_API_KEY, DD_SITE, and the agent's provider key(s)
# (e.g. OPENAI_API_KEY). NOT DD_APP_KEY — this replays a plain trace, not an Experiment.
import argparse
import asyncio
import json
import os
import sys
from dotenv import load_dotenv
# override=True makes the project's .env authoritative: a developer shell commonly exports ambient DD_*
# vars (e.g. for the datadog-llmo MCP, often a different org) that would otherwise win over .env and send
# the trace to the wrong org. MUST precede the ddtrace import. (It won't clobber the caller's DD_TAGS
# unless .env also sets DD_TAGS.)
load_dotenv(override=True)

from ddtrace.llmobs import LLMObs

# Import each instrumented entrypoint:
from {{MODULE}} import {{ENTRYPOINT_FN}}   # , {{ANOTHER_ENTRYPOINT_FN}}, ...

ML_APP = os.environ.get("DD_LLMOBS_ML_APP", "{{ML_APP}}")
# Local replays emit under "<ml_app>-local" (idempotent) so they never pollute the production ml_app.
if not ML_APP.endswith("-local"):
    ML_APP = ML_APP + "-local"
# Interlock: refuse to emit under a non-isolated ml_app. NOTE this only guards the init-level setting — if
# the app sets ml_app per span/call it overrides this (see the pre-flight in references/details.md).
if not ML_APP.endswith("-local"):
    raise SystemExit(f"[replay] refusing to run: ml_app {ML_APP!r} is not isolated (must end in -local)")

# Export mode: agentless is the right default for a LOCAL replay (ships LLM Obs spans straight to Datadog,
# no Agent needed). If this app is instead wired to a local Agent sidecar, set DD_LLMOBS_AGENTLESS_ENABLED=0.
# Benign gotcha: with agentless on, the APM tracer may still dial localhost:8126 and log
# "ERROR: lost N traces ... connection refused". That is HARMLESS — the LLM Obs spans ship independently and
# arrive fine — so don't mistake it for a failed replay.
_agentless = os.environ.get("DD_LLMOBS_AGENTLESS_ENABLED", "1").lower() not in ("0", "false", "no")
LLMObs.enable(ml_app=ML_APP, agentless_enabled=_agentless)


# Dispatch table — ONE entry per execution type, keyed by the replay_entrypoint id the app annotates.
# Each entry: the function to call + whether it's async. The runner calls fn(**input_data); the trace it
# emits is what the skill diffs, so no return-value extraction is needed here.
ENTRYPOINTS = {
    # "{{REPLAY_ENTRYPOINT_ID}}": {"fn": {{ENTRYPOINT_FN}}, "is_async": True},
}


def _run_entrypoint(spec, input_data):
    fn, is_async = spec["fn"], spec.get("is_async", False)
    return asyncio.run(fn(**input_data)) if is_async else fn(**input_data)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--entrypoint", required=True, help="replay_entrypoint id (key in ENTRYPOINTS)")
    ap.add_argument("--input-file", required=True, help="path to a JSON file of entrypoint kwargs")
    args = ap.parse_args()

    if args.entrypoint not in ENTRYPOINTS:
        print(json.dumps({"error": f"unknown entrypoint {args.entrypoint!r}; known: {sorted(ENTRYPOINTS)}"}))
        sys.exit(2)
    with open(args.input_file) as f:
        input_data = json.load(f)

    # Run the entrypoint DIRECTLY — no wrapper span — so the replay trace is structurally identical to a
    # normal run. The correlation marker rides along as a span tag via DD_TAGS (set by the caller).
    # Flush in `finally` so a FAILED replay still emits its partial trace — otherwise the error path leaves
    # nothing to diff, which is worse than a visible failure.
    status, error = "done", None
    try:
        _run_entrypoint(ENTRYPOINTS[args.entrypoint], input_data)
    except Exception as exc:  # noqa: BLE001 — normal failures → report + exit 1
        status, error = "error", repr(exc)
    finally:
        LLMObs.flush()  # send the trace (even a partial one) before we exit
        # Print in `finally` so the caller always gets ml_app to poll — even if the entrypoint raised
        # SystemExit/KeyboardInterrupt (which skip `except Exception` but still run `finally`, then propagate).
        print(json.dumps({"status": status, "entrypoint": args.entrypoint, "ml_app": ML_APP, "error": error}))
    if status == "error":
        sys.exit(1)


if __name__ == "__main__":
    main()