import os
from pprint import pprint

from trio import ServiceClient


def print_block(title: str, payload):
    print(f"\n===== {title} =====")
    pprint(payload)


def first_or_none(items):
    return items[0] if items else None


domain = os.getenv("SERVICE_DOMAIN", "117.50.44.26:3000")
allow_delete = os.getenv("RUN_DELETE_REST_TESTS", "false").lower() == "true"

# 创建 rest 客户端，用于获取模型列表和模型下载链接
client = ServiceClient(domain)
rest_client = client.create_rest_client()

# 获取用户信息
user_info = rest_client.get_user_info()
print_block("get_user_info", user_info)

weights = rest_client.list_weights()
print_block("list_weights", weights)

legacy_checkpoints = weights.get("data", [])
legacy_checkpoint = first_or_none(legacy_checkpoints)
if legacy_checkpoint:
    legacy_archive = rest_client.get_archive_url(legacy_checkpoint["cuid"])
    print_block("get_archive_url", legacy_archive)
else:
    print("\n===== get_archive_url =====")
    print("skip: no legacy checkpoint found")

# 2) new restclient-compatible methods
training_runs = rest_client.list_training_runs(limit=20, offset=0)
print_block("list_training_runs", training_runs)

training_run = first_or_none(training_runs.get("training_runs", []))
if training_run:
    training_run_id = training_run["training_run_id"]

    training_run_detail = rest_client.get_training_run(training_run_id)
    print_block("get_training_run", training_run_detail)

    run_checkpoints = rest_client.list_checkpoints(training_run_id)
    print_block("list_checkpoints", run_checkpoints)

    run_checkpoint = first_or_none(run_checkpoints.get("checkpoints", []))
    if run_checkpoint:
        checkpoint_id = run_checkpoint["checkpoint_id"]

        checkpoint_archive = rest_client.get_checkpoint_archive_url(
            training_run_id, checkpoint_id
        )
        print_block("get_checkpoint_archive_url", checkpoint_archive)

        if allow_delete:
            rest_client.delete_checkpoint(training_run_id, checkpoint_id)
            print("\n===== delete_checkpoint =====")
            print(
                f"deleted checkpoint {checkpoint_id} under training run {training_run_id}"
            )
        else:
            print("\n===== delete_checkpoint =====")
            print("skip: set RUN_DELETE_REST_TESTS=true to enable deletion")
    else:
        print("\n===== run checkpoint actions =====")
        print("skip: no checkpoint found under the first training run")
else:
    print("\n===== training run actions =====")
    print("skip: no training run found")

user_checkpoints = rest_client.list_user_checkpoints(limit=100, offset=0)
print_block("list_user_checkpoints", user_checkpoints)

sessions = rest_client.list_sessions(limit=20, offset=0)
print_block("list_sessions", sessions)

session = first_or_none(sessions.get("sessions", []))
if session:
    session_id = session["session_id"]

    session_detail = rest_client.get_session(session_id)
    print_block("get_session", session_detail)

    sampler_id = first_or_none(session_detail.get("sampler_ids", []))
    if sampler_id:
        sampler_detail = rest_client.get_sampler(sampler_id)
        print_block("get_sampler", sampler_detail)
    else:
        print("\n===== get_sampler =====")
        print("skip: no sampler id found in the first session")
else:
    print("\n===== session actions =====")
    print("skip: no session found")
