#!/usr/bin/env python3
"""
O11Y S3 Data Explorer - Navigate and query O11Y S3 data.

Usage:
    python o11y-s3-explore.py list-clusters <region>
    python o11y-s3-explore.py cluster-dirs <cluster_id> [region]
    python o11y-s3-explore.py log-stats <cluster_id> <hour_YYYYMMDDHH> [region]
    python o11y-s3-explore.py list-logs <cluster_id> <start_YYYYMMDDHH> <end_YYYYMMDDHH> [region]
    python o11y-s3-explore.py devtier-logs <env> <tenant_id> <project_id> <cluster_id> [region]
    python o11y-s3-explore.py list-deltalake <tenant_id> [region]

Requires AWS credentials in environment (use o11y-data.sh o11y_auth first).
"""
import sys
import os
import tempfile
import sqlite3


def get_s3_client(region='us-west-2'):
    import boto3
    return boto3.client('s3', region_name=region)


def bucket_name(region):
    return f'o11y-prod-shared-{region}'


def cmd_list_clusters(region='us-west-2'):
    """List cluster IDs under diagnosis/data/"""
    s3 = get_s3_client(region)
    paginator = s3.get_paginator('list_objects_v2')
    count = 0
    for page in paginator.paginate(Bucket=bucket_name(region), Prefix='diagnosis/data/', Delimiter='/'):
        for p in page.get('CommonPrefixes', []):
            cid = p['Prefix'].split('/')[-2]
            print(cid)
            count += 1
    print(f'\nTotal: {count} clusters', file=sys.stderr)


def cmd_cluster_dirs(cluster_id, region='us-west-2'):
    """Show directories available for a given cluster."""
    s3 = get_s3_client(region)
    prefix = f'diagnosis/data/{cluster_id}/'
    resp = s3.list_objects_v2(Bucket=bucket_name(region), Prefix=prefix, Delimiter='/')
    dirs = [p['Prefix'].replace(prefix, '').rstrip('/') for p in resp.get('CommonPrefixes', [])]
    for d in dirs:
        print(d)


def cmd_log_stats(cluster_id, hour, region='us-west-2'):
    """Download and display log-stats for a specific hour."""
    s3 = get_s3_client(region)
    key = f'diagnosis/data/{cluster_id}/log-stats/1h/{hour}.db'
    
    with tempfile.NamedTemporaryFile(suffix='.db', delete=False) as tmp:
        try:
            s3.download_file(bucket_name(region), key, tmp.name)
        except Exception as e:
            print(f'Error downloading {key}: {e}', file=sys.stderr)
            return
    
    try:
        conn = sqlite3.connect(tmp.name)
        cur = conn.cursor()
        cur.execute('SELECT component, level, position, msg_digest, count, sample FROM log_stats ORDER BY count DESC')
        rows = cur.fetchall()
        
        print(f'{"COMPONENT":<10} {"LEVEL":<6} {"COUNT":>8}  {"POSITION":<40} {"MESSAGE"}')
        print('-' * 120)
        for row in rows:
            component, level, position, msg_digest, count, sample = row
            print(f'{component:<10} {level:<6} {count:>8}  {position[:40]:<40} {msg_digest[:60]}')
        
        print(f'\nTotal: {len(rows)} entries')
        conn.close()
    finally:
        os.unlink(tmp.name)


def cmd_list_logs(cluster_id, start, end, region='us-west-2'):
    """List merged log files for a time range (YYYYMMDDHH format)."""
    s3 = get_s3_client(region)
    bkt = bucket_name(region)
    prefix = f'diagnosis/data/{cluster_id}/merged-logs/'
    
    resp = s3.list_objects_v2(Bucket=bkt, Prefix=prefix, Delimiter='/')
    dirs = [p['Prefix'].replace(prefix, '').rstrip('/') for p in resp.get('CommonPrefixes', [])]
    
    components = [d for d in dirs if not d[0].isdigit()]
    hour_dirs = [d for d in dirs if d[0].isdigit() and start <= d <= end]
    
    for comp in components:
        resp = s3.list_objects_v2(Bucket=bkt, Prefix=f'{prefix}{comp}/', MaxKeys=100)
        for obj in resp.get('Contents', []):
            fname = os.path.basename(obj['Key']).replace('.log', '')
            if start <= fname <= end:
                print(f'{obj["Key"]}  ({obj["Size"]} bytes)')

    for hdir in hour_dirs:
        resp = s3.list_objects_v2(Bucket=bkt, Prefix=f'{prefix}{hdir}/', MaxKeys=100)
        for obj in resp.get('Contents', []):
            print(f'{obj["Key"]}  ({obj["Size"]} bytes)')


def cmd_devtier_logs(env, tenant_id, project_id, cluster_id, region='us-west-2'):
    """List log directories for essential/starter clusters."""
    s3 = get_s3_client(region)
    bkt = bucket_name(region)
    
    for base in [
        f'0/{tenant_id}/{project_id}/{cluster_id}/logs/',
        f'0/{env}-devtier/serverless/{tenant_id}/{project_id}/{cluster_id}/logs/'
    ]:
        resp = s3.list_objects_v2(Bucket=bkt, Prefix=base, Delimiter='/', MaxKeys=50)
        prefixes = resp.get('CommonPrefixes', [])
        if prefixes:
            print(f'Found at: {base}')
            for p in prefixes:
                dirname = p['Prefix'].replace(base, '').rstrip('/')
                print(f'  {dirname}/')
            return
    
    print('No logs found for the given cluster parameters.', file=sys.stderr)


def cmd_list_deltalake(tenant_id, region='us-west-2'):
    """List Delta Lake tables for a tenant."""
    s3 = get_s3_client(region)
    bkt = bucket_name(region)
    prefix = f'deltalake/{tenant_id}/'
    
    resp = s3.list_objects_v2(Bucket=bkt, Prefix=prefix, Delimiter='/')
    for p in resp.get('CommonPrefixes', []):
        oid = p['Prefix'].replace(prefix, '').rstrip('/')
        r2 = s3.list_objects_v2(Bucket=bkt, Prefix=p['Prefix'], Delimiter='/')
        tables = [pp['Prefix'].split('/')[-2] for pp in r2.get('CommonPrefixes', [])]
        print(f'{oid}: {", ".join(tables)}')


if __name__ == '__main__':
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)
    
    cmd = sys.argv[1]
    args = sys.argv[2:]
    
    commands = {
        'list-clusters': cmd_list_clusters,
        'cluster-dirs': cmd_cluster_dirs,
        'log-stats': cmd_log_stats,
        'list-logs': cmd_list_logs,
        'devtier-logs': cmd_devtier_logs,
        'list-deltalake': cmd_list_deltalake,
    }
    
    if cmd in commands:
        try:
            commands[cmd](*args)
        except TypeError as e:
            print(f'Error: {e}', file=sys.stderr)
            print(f'Usage: python {sys.argv[0]} {cmd} ...', file=sys.stderr)
    else:
        print(f'Unknown command: {cmd}', file=sys.stderr)
        print(f'Available: {", ".join(commands.keys())}', file=sys.stderr)