#!/usr/bin/env python
# (C) 2026 Cadence Design Systems, Inc. (Cadence)
# All rights reserved.
# TERMS FOR USE OF SAMPLE CODE The software below ("Sample Code") is
# provided to current licensees or subscribers of Cadence products or
# SaaS offerings (each a "Customer").
# Customer is hereby permitted to use, copy, and modify the Sample Code,
# subject to these terms. Cadence claims no rights to Customer's
# modifications. Modification of Sample Code is at Customer's sole and
# exclusive risk. Sample Code may require Customer to have a then
# current license or subscription to the applicable Cadence offering.
# THE SAMPLE CODE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED.  CADENCE DISCLAIMS ALL WARRANTIES, INCLUDING, BUT
# NOT LIMITED TO, WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
# PARTICULAR PURPOSE AND NONINFRINGEMENT. In no event shall Cadence be
# liable for any damages or liability in connection with the Sample Code
# or its use.
"""List and fetch MD trajectories from an Orion shard collection.

    python oemdaffinity_list_collection.py 410438 --profile staging --grep ejm_43
    python oemdaffinity_list_collection.py 410438 --profile staging --grep ejm_43 --download traj
    python oemdaffinity_list_collection.py 410438 --profile staging --kind all
"""

import os
import sys
import argparse

from orionclient.session import OrionSession, get_profile_config
from orionclient.types import ShardCollection


def leg_of(name):
    if "_bn" in name:
        return "bound"
    if "_unb" in name:
        return "unbound"
    return ""


def human(n):
    if n is None:
        return "?"
    for u in ("B", "KB", "MB", "GB"):
        if n < 1024:
            return f"{n:.1f}{u}"
        n /= 1024
    return f"{n:.1f}TB"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("collection_id", type=int)
    ap.add_argument("--profile", default="default")
    ap.add_argument("--grep")
    ap.add_argument("--kind", choices=["trajectory", "all"], default="trajectory")
    ap.add_argument("--leg", choices=["bound", "unbound"])
    ap.add_argument("--download", metavar="DIR")
    args = ap.parse_args()

    session = OrionSession(config=get_profile_config(profile=args.profile))
    collection = session.get_resource(ShardCollection, args.collection_id)
    print(f"Collection {collection.id}: {collection.name!r}  state={collection.state}\n")

    shards = list(collection.list_shards())
    truncated = [s for s in shards if (s.name or "").endswith("...")]

    if args.kind == "trajectory":
        shards = [s for s in shards if (s.name or "").endswith("_traj.tar.gz")]
    if args.leg:
        shards = [s for s in shards if leg_of(s.name or "") == args.leg]
    if args.grep:
        shards = [s for s in shards if args.grep.lower() in (s.name or "").lower()]

    print(f"{'ID':>19}  {'LEG':<8} {'SIZE':>9}  NAME")
    print("-" * 120)
    for s in sorted(shards, key=lambda s: s.name or ""):
        print(f"{s.id:>19}  {leg_of(s.name or ''):<8} {human(getattr(s, 'size', None)):>9}  {s.name}")
        if getattr(s, "metadata", None):
            print(f"{'':>19}  meta: {dict(s.metadata)}")
    print("-" * 120)
    print(f"{len(shards)} shards, total {human(sum(getattr(s, 'size', None) or 0 for s in shards))}")

    if truncated:
        print(f"\nWARNING: {len(truncated)} shard name(s) truncated at 64 chars; "
              f"'_traj' may be lost, so --kind trajectory can under-report. "
              f"Check metadata['full_name'] on: {[s.id for s in truncated]}", file=sys.stderr)

    if args.download:
        os.makedirs(args.download, exist_ok=True)
        for s in shards:
            dest = os.path.join(args.download, s.name or f"shard_{s.id}")
            print(f"{s.id} -> {dest}  ({human(getattr(s, 'size', None))})")
            s.download_to_file(dest)
    return 0


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