#!/usr/bin/env python3
"""Create portable copies of locally derived manifests and reports."""

from __future__ import annotations

import argparse
import json
import shutil
from pathlib import Path


ACTION_MANIFESTS = {
    "autoeval_full.jsonl": "autoeval",
    "libero_full.jsonl": "libero_success",
    "roboarena_full_refined.jsonl": "roboarena",
    "racer_pilot.jsonl": "racer_pilot",
}
ACTION_REPORTS = (
    "summary.json",
    "autoeval_full_validation.json",
    "libero_full_validation.json",
    "roboarena_full_refined_validation.json",
    "racer_pilot_validation.json",
)


def portable(value, experiment_root: Path):
    if isinstance(value, dict):
        return {key: portable(item, experiment_root) for key, item in value.items()}
    if isinstance(value, list):
        return [portable(item, experiment_root) for item in value]
    if not isinstance(value, str):
        return value
    if "/action_labels/stage2_action_sidecar_v1/" in value:
        return "action_sidecars/" + value.split("/action_labels/stage2_action_sidecar_v1/", 1)[1]
    racer_marker = "/artifacts/action_recovery_v1/racer_pilot/"
    if racer_marker in value:
        return "action_sidecars/racer_pilot/" + value.split(racer_marker, 1)[1]
    if "/action_sources/" in value:
        return "upstream_raw/" + value.split("/action_sources/", 1)[1]
    root_text = str(experiment_root.resolve()) + "/"
    if value.startswith(root_text):
        return value[len(root_text):]
    return value


def rewrite_json(source: Path, target: Path, experiment_root: Path) -> None:
    payload = portable(json.loads(source.read_text()), experiment_root)
    target.parent.mkdir(parents=True, exist_ok=True)
    target.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")


def rewrite_jsonl(source: Path, target: Path, experiment_root: Path, source_name: str) -> None:
    target.parent.mkdir(parents=True, exist_ok=True)
    output = []
    for line in source.read_text().splitlines():
        if not line:
            continue
        row = portable(json.loads(line), experiment_root)
        if "sidecar_path" in row:
            row["sidecar_path"] = f"action_sidecars/{source_name}/trajectory_{row['stage2_trajectory_id']}.npz"
        output.append(json.dumps(row, sort_keys=True) + "\n")
    target.write_text("".join(output))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--experiment-root", type=Path, required=True)
    parser.add_argument("--dataset-metadata-root", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    args = parser.parse_args()
    experiment = args.experiment_root.resolve()
    output = args.output_root.resolve()

    manifests = experiment / "manifests"
    shutil.copy2(manifests / "phase2_remediation_pairs_v1.jsonl", output / "manifests/phase2_remediation_pairs_v1.jsonl")
    rewrite_json(manifests / "phase2_remediation_pairs_v1_summary.json", output / "manifests/phase2_remediation_pairs_v1_summary.json", experiment)
    action_source = manifests / "action_recovery_v1"
    for name, source_name in ACTION_MANIFESTS.items():
        rewrite_jsonl(action_source / name, output / "derived/action_manifests" / name, experiment, source_name)
    report_source = experiment / "reports/action_recovery_v1"
    for name in ACTION_REPORTS:
        rewrite_json(report_source / name, output / "derived/action_reports" / name, experiment)

    inventory_source = args.dataset_metadata_root / "phase21a_data_inventory.json"
    inventory = json.loads(inventory_source.read_text())
    inventory["manifest"] = "manifests/phase2_remediation_pairs_v1.jsonl"
    inventory["output_dir"] = "."
    (output / "manifests/stage2_subset_inventory.json").write_text(
        json.dumps(inventory, indent=2, sort_keys=True) + "\n"
    )
    shutil.copy2(args.dataset_metadata_root / "SHA256SUMS", output / "manifests/stage2_subset_SHA256SUMS")
    print(json.dumps({"status": "complete", "output_root": str(output)}, indent=2))


if __name__ == "__main__":
    main()