# Created by BVH, Apr 2026.
# A4D-specific run logging helpers.

import sys

from rich import print
from rich.panel import Panel


def get_latest_s3_checkpoint(s3_path):
    '''Find the latest model checkpoint on S3 for a given run. Returns S3 path or empty string.
    :param s3_path: full S3 path for the run (e.g. s3://bucket/prefix/a4d2/dgx/...).
    '''
    if not s3_path:
        return ''
    try:
        import boto3
        # Parse bucket and prefix from the full S3 URI
        stripped = s3_path.replace('s3://', '')
        bucket = stripped.split('/')[0]
        run_prefix = '/'.join(stripped.split('/')[1:])
        s3 = boto3.client('s3')
        for suffix in ('/model/', '/checkpoints/model/'):
            resp = s3.list_objects_v2(Bucket=bucket, Prefix=run_prefix + suffix)
            files = sorted([o['Key'] for o in resp.get('Contents', [])])
            if files:
                return f's3://{bucket}/{files[-1]}'
    except Exception:
        pass
    return ''


def print_run_info(phase, train_step, local_path, s3_path):
    '''Print run paths after visuals are saved. Rank 0 only.
    :param local_path: absolute local directory (e.g. /data/a4d2/dgx/...).
    :param s3_path: full S3 path for the run, or empty string if S3 sync disabled.
    '''
    import torch
    if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0:
        return
    latest_ckpt = get_latest_s3_checkpoint(s3_path)
    anyfile_base = 'https://anyfile.us-west-2.awsinternal.tri.global/browse'
    anyfile_link = f'{anyfile_base}/{s3_path}' if s3_path else ''
    body = (
        f'[white]Command:[/]      [dim]{" ".join(sys.argv)}[/]\n'
        f'[white]Local path:[/]   [green]{local_path}[/]\n'
        f'[white]Visuals:[/]      [green]{local_path}/visuals/[/]\n'
        f'[white]S3 path:[/]      [blue]{s3_path or "(disabled)"}[/]\n'
        f'[white]AnyFile:[/]      [cyan]{anyfile_link or "(disabled)"}[/]\n'
        f'[white]Latest ckpt:[/]  [magenta]{latest_ckpt or "(none yet)"}[/]'
    )
    print(Panel(body, title=f'[bold yellow]RUN INFO ({phase} s{train_step})[/]',
                border_style='yellow', expand=False))