"""Export contimp-app traces from LangFuse as observation-level JSONL.

Produces rows shaped like LangFuse's Tracing-tab batch export (one JSON object
per observation, trace-level fields duplicated, input/output JSON-encoded) —
the format the oumi platform's LangFuse importer consumes.

    uv run python scripts/export_traces.py --tag pr-area --limit 50 --out /tmp/export.jsonl

Auth via LANGFUSE_PUBLIC_KEY / LANGFUSE_SECRET_KEY / LANGFUSE_BASE_URL env vars.
"""

import argparse
import base64
import json
import os
import urllib.parse
import urllib.request


def api_get(path: str, params: dict | None = None) -> dict:
    base = os.environ.get("LANGFUSE_BASE_URL", "https://cloud.langfuse.com").rstrip("/")
    auth = base64.b64encode(
        f"{os.environ['LANGFUSE_PUBLIC_KEY']}:{os.environ['LANGFUSE_SECRET_KEY']}".encode()
    ).decode()
    query = f"?{urllib.parse.urlencode(params)}" if params else ""
    req = urllib.request.Request(
        f"{base}/api/public{path}{query}", headers={"Authorization": f"Basic {auth}"}
    )
    return json.loads(urllib.request.urlopen(req, timeout=60).read())


def encode(value: object) -> str | None:
    if value is None:
        return None
    return value if isinstance(value, str) else json.dumps(value, ensure_ascii=False)


def export_trace(trace_id: str) -> list[dict]:
    trace = api_get(f"/traces/{trace_id}")
    rows = []
    for obs in trace.get("observations", []):
        rows.append({
            "id": obs.get("id"),
            "traceId": trace["id"],
            "traceName": trace.get("name"),
            "sessionId": trace.get("sessionId"),
            "userId": trace.get("userId"),
            "traceTags": trace.get("tags"),
            "traceMetadata": trace.get("metadata"),
            "traceScores": [
                {"name": s.get("name"), "value": s.get("value"), "comment": s.get("comment")}
                for s in trace.get("scores", [])
            ],
            "type": obs.get("type"),
            "name": obs.get("name"),
            "startTime": obs.get("startTime"),
            "endTime": obs.get("endTime"),
            "parentObservationId": obs.get("parentObservationId"),
            "model": obs.get("model"),
            "input": encode(obs.get("input")),
            "output": encode(obs.get("output")),
            "metadata": obs.get("metadata"),
        })
    return rows


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--tag", help="filter traces by tag (e.g. a task id)")
    parser.add_argument("--limit", type=int, default=100)
    parser.add_argument("--out", default="/tmp/langfuse_export.jsonl")
    args = parser.parse_args()

    params: dict = {"limit": min(args.limit, 100), "page": 1}
    if args.tag:
        params["tags"] = args.tag
    trace_ids: list[str] = []
    while len(trace_ids) < args.limit:
        page = api_get("/traces", params)
        trace_ids += [t["id"] for t in page["data"]]
        if params["page"] >= page["meta"]["totalPages"]:
            break
        params["page"] += 1
    trace_ids = trace_ids[: args.limit]

    with open(args.out, "w") as f:
        n_rows = 0
        for trace_id in trace_ids:
            for row in export_trace(trace_id):
                f.write(json.dumps(row, ensure_ascii=False) + "\n")
                n_rows += 1
    print(f"exported {len(trace_ids)} traces ({n_rows} observations) to {args.out}")


if __name__ == "__main__":
    main()