#!/usr/bin/env python3
"""Copy Braintrust projects, datasets and experiments into Oodle.

The script reads from the Braintrust REST API and writes to the
Braintrust-compatible API of an Oodle instance. It needs only the
Python standard library (3.9 or later).

    python braintrust_import.py \\
        --source-api-key "$BRAINTRUST_API_KEY" \\
        --dest-url "https://<API_DOMAIN>/v1/api/instance/<OODLE_INSTANCE>/braintrust" \\
        --dest-api-key "$OODLE_API_KEY" \\
        --project enterpret-agent-local --project wisdom-mcp-staging

A run is safe to repeat. Projects, datasets and experiments are
found by name, and rows keep their Braintrust ids, so a second
run updates what the first one wrote and adds no copies.

Production logs are not copied: they are traces, and new traces
reach Oodle when the SDK points at it.
"""

from __future__ import annotations

import argparse
import json
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

DEFAULT_SOURCE_URL = "https://api.braintrust.dev"
PAGE = 100
FETCH_PAGE = 1000
WRITE_BATCH = 200
RETRIES = 5
TIMEOUT_S = 120

# Fields of a fetched row that describe the row in the source and
# are not data to copy.
SKIP_FIELDS = {
    "_xact_id",
    "_pagination_key",
    "_object_delete",
    "_is_merge",
    "created",
    "project_id",
    "experiment_id",
    "dataset_id",
    "log_id",
    "org_id",
    "is_root",
}


class Client:
    """A JSON client with retries for one API base URL."""

    def __init__(self, base: str, api_key: str):
        self.base = base.rstrip("/")
        self.api_key = api_key

    def request(self, method: str, path: str, body=None, query=None):
        url = self.base + path
        if query:
            url += "?" + urllib.parse.urlencode(
                {k: v for k, v in query.items() if v is not None}, doseq=True
            )
        data = None if body is None else json.dumps(body).encode()
        delay = 1.0
        for attempt in range(RETRIES):
            req = urllib.request.Request(url, data=data, method=method)
            req.add_header("Authorization", "Bearer " + self.api_key)
            req.add_header("Content-Type", "application/json")
            try:
                with urllib.request.urlopen(req, timeout=TIMEOUT_S) as resp:
                    raw = resp.read()
                    return json.loads(raw) if raw else {}
            except urllib.error.HTTPError as err:
                text = err.read().decode(errors="replace")[:500]
                retryable = err.code == 429 or err.code >= 500
                if not retryable or attempt == RETRIES - 1:
                    raise RuntimeError(f"{method} {url}: HTTP {err.code}: {text}") from None
            except urllib.error.URLError as err:
                if attempt == RETRIES - 1:
                    raise RuntimeError(f"{method} {url}: {err.reason}") from None
            time.sleep(delay)
            delay = min(delay * 2, 30)
        raise RuntimeError(f"{method} {url}: no response")


def list_objects(client: Client, path: str, query: dict):
    """Yield every object of a list endpoint, page by page.

    A page that adds no new id ends the listing, so a server that
    does not page by `starting_after` cannot make this loop
    forever.
    """
    seen = set()
    after = None
    while True:
        q = dict(query, limit=PAGE)
        if after:
            q["starting_after"] = after
        objects = client.request("GET", path, query=q).get("objects") or []
        new = [o for o in objects if o.get("id") not in seen]
        for o in new:
            seen.add(o["id"])
            yield o
        if len(objects) < PAGE or not new:
            return
        after = new[-1]["id"]


def fetch_rows(client: Client, path: str):
    """Yield every row of an object through its fetch endpoint."""
    cursor = None
    while True:
        body = {"limit": FETCH_PAGE}
        if cursor:
            body["cursor"] = cursor
        resp = client.request("POST", path, body=body)
        for event in resp.get("events") or []:
            yield event
        cursor = resp.get("cursor")
        if not cursor:
            return


def clean_row(event: dict) -> dict:
    return {k: v for k, v in event.items() if k not in SKIP_FIELDS and v is not None}


def write_rows(dest: Client, rows: list, dry_run: bool) -> int:
    if dry_run:
        return len(rows)
    for i in range(0, len(rows), WRITE_BATCH):
        dest.request("POST", "/logs3", body={"rows": rows[i : i + WRITE_BATCH], "api_version": 2})
    return len(rows)


def import_project(source: Client, dest: Client, project: dict, dry_run: bool, log) -> dict:
    name = project["name"]
    counts = {"datasets": 0, "dataset_rows": 0, "experiments": 0, "experiment_rows": 0}
    if dry_run:
        dest_project_id = "dry-run"
    else:
        dest_project_id = dest.request("POST", "/v1/project", body={"name": name})["id"]
    log(f"project {name}")

    dataset_ids = {}
    for ds in list_objects(source, "/v1/dataset", {"project_id": project["id"]}):
        if dry_run:
            dest_id = "dry-run"
        else:
            dest_id = dest.request(
                "POST",
                "/v1/dataset",
                body={
                    "project_id": dest_project_id,
                    "name": ds["name"],
                    "description": ds.get("description"),
                    "metadata": ds.get("metadata"),
                },
            )["id"]
        dataset_ids[ds["id"]] = dest_id
        rows = []
        for event in fetch_rows(source, f"/v1/dataset/{ds['id']}/fetch"):
            row = clean_row(event)
            row = {k: row[k] for k in ("id", "input", "expected", "metadata", "tags") if k in row}
            row["dataset_id"] = dest_id
            rows.append(row)
        counts["datasets"] += 1
        counts["dataset_rows"] += write_rows(dest, rows, dry_run)
        log(f"  dataset {ds['name']}: {len(rows)} rows")

    experiments = list(list_objects(source, "/v1/experiment", {"project_id": project["id"]}))
    experiments.sort(key=lambda e: e.get("created") or "")
    experiment_ids = {}
    for exp in experiments:
        body = {
            "project_id": dest_project_id,
            "name": exp["name"],
            "description": exp.get("description"),
            "repo_info": exp.get("repo_info"),
            "metadata": exp.get("metadata"),
            "tags": exp.get("tags"),
            "ensure_new": False,
        }
        if exp.get("dataset_id") in dataset_ids:
            body["dataset_id"] = dataset_ids[exp["dataset_id"]]
        dest_id = "dry-run" if dry_run else dest.request("POST", "/v1/experiment", body=body)["id"]
        experiment_ids[exp["id"]] = dest_id
        rows = []
        for event in fetch_rows(source, f"/v1/experiment/{exp['id']}/fetch"):
            row = clean_row(event)
            row["experiment_id"] = dest_id
            origin = row.get("origin")
            if isinstance(origin, dict) and origin.get("object_id") in dataset_ids:
                row["origin"] = dict(origin, object_id=dataset_ids[origin["object_id"]])
            rows.append(row)
        counts["experiments"] += 1
        counts["experiment_rows"] += write_rows(dest, rows, dry_run)
        log(f"  experiment {exp['name']}: {len(rows)} rows")

    # Base links are set after every experiment exists: two
    # experiments can share a creation time, so the list order
    # does not say which one came first.
    for exp in experiments:
        base = exp.get("base_exp_id")
        if base in experiment_ids and not dry_run:
            dest.request(
                "PATCH",
                f"/v1/experiment/{experiment_ids[exp['id']]}",
                body={"base_exp_id": experiment_ids[base]},
            )
    return counts


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    parser.add_argument("--source-url", default=DEFAULT_SOURCE_URL, help="Braintrust API URL")
    parser.add_argument("--source-api-key", required=True, help="Braintrust API key")
    parser.add_argument("--dest-url", required=True, help="Oodle Braintrust-compatible API URL")
    parser.add_argument("--dest-api-key", required=True, help="Oodle API key")
    parser.add_argument(
        "--project", action="append", default=[], help="project name to copy; repeat it; default all"
    )
    parser.add_argument("--dry-run", action="store_true", help="read and count, write nothing")
    parser.add_argument("--quiet", action="store_true")
    args = parser.parse_args(argv)

    source = Client(args.source_url, args.source_api_key)
    dest = Client(args.dest_url, args.dest_api_key)

    def log(msg):
        if not args.quiet:
            print(msg, file=sys.stderr)

    wanted = set(args.project)
    totals = {"projects": 0, "datasets": 0, "dataset_rows": 0, "experiments": 0, "experiment_rows": 0}
    found = set()
    for project in list_objects(source, "/v1/project", {}):
        if wanted and project["name"] not in wanted:
            continue
        found.add(project["name"])
        counts = import_project(source, dest, project, args.dry_run, log)
        totals["projects"] += 1
        for k, v in counts.items():
            totals[k] += v
    missing = wanted - found
    if missing:
        log("projects not found in the source: " + ", ".join(sorted(missing)))
    print(json.dumps(totals))
    return 1 if missing else 0


if __name__ == "__main__":
    sys.exit(main())