#!/usr/bin/env python3
"""
Clean up small S3 log files that have already been merged, and collect compression stats.

Small-file path:  0/{tenantid}/{projectid}/{clusterid}/logs/{component}/{timestamp}-{uuid}.log.gz
Merged path:      diagnosis/data/{clusterid}/merged-logs/{component}/{YYYYMMDDHH}.log

Logic: list small files from a start time; delete if the corresponding hour already has a merged
file; otherwise record file count/size per cluster-component.
Minimize S3 API calls: list by 3-digit coarse prefixes (e.g. 175, 176) so each base only needs a
few lists; probe with 2-digit prefixes (e.g. 17) using MaxKeys=1 and skip all lists under that
2-digit base when empty.
Also count small-file number and total size to estimate post-merge compression.

# Stats only + preview keys that would be deleted (no delete)
python cleanup.py \
  --bucket o11y-prod-shared-ap-southeast-3 \
  --start-time "2026-02-20" \
  --end-time "2026-02-21" \
  --region ap-southeast-3 \
  --dry-run

# List and stats only; do not delete or print keys to delete
python scripts/cleanup_duplicate_s3_logs.py \
  --bucket o11y-prod-shared-us-east-1 \
  --start-time "2026-02-20" \
  --skip-delete

# Bounded time range with real deletes (streaming; delete every 1000 keys to limit memory)
python scripts/cleanup_duplicate_s3_logs.py \
  --bucket o11y-prod-shared-us-east-1 \
  --start-time "2026-02-20" \
  --end-time "2026-02-26" \
  --region us-east-1

# Phase 2: skip dirs (any path segment matching these names is excluded from list/delete)
python scripts/cleanup_duplicate_s3_logs.py ... --exclude-dirs prod-devtier dev-devtier 30008 30073

# Progress file: after crash, --resume continues from last position without full rescan
python scripts/cleanup_duplicate_s3_logs.py \
  --bucket o11y-prod-shared-us-east-1 \
  --start-time "2026-02-20" \
  --end-time "2026-02-26" \
  --progress-file ./cleanup_progress.json
# Resume: add --resume
python scripts/cleanup_duplicate_s3_logs.py ... --progress-file ./cleanup_progress.json --resume
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any, Callable, Iterator, Optional

try:
    import boto3
    from botocore.exceptions import ClientError
except ImportError:
    boto3 = None
    ClientError = None

# Small-file key format: 0/{tenant}/{project}/{cluster}/logs/{component}/{timestamp}-{uuid}.log.gz
SMALL_KEY_PATTERN = re.compile(
    r"^0/([^/]+)/([^/]+)/([^/]+)/logs/([^/]+)/(\d+)-[^/]+\.log\.gz$"
)
# Merged path format (same as diagnosis-compute): diagnosis/data/{cluster}/merged-logs/{component}/{YYYYMMDDHH}.log
MERGED_KEY_PATTERN = re.compile(
    r"^diagnosis/data/([^/]+)/merged-logs/([^/]+)/(\d{10})\.log$"
)
# Also supports: diagnosis/data/{cluster}/merged-logs/{datetime}/{component}/xxx.log
MERGED_KEY_PATTERN_ALT = re.compile(
    r"^diagnosis/data/([^/]+)/merged-logs/(\d{10})/([^/]+)/[^/]+\.log$"
)

LIST_PAGE_SIZE = 1000  # list_objects_v2 page size; S3 max is 1000
DELETE_BATCH_SIZE = 1000  # delete_objects batch size; S3 max is 1000
# Use leading timestamp digits as prefix when listing small files to avoid a full scan of 0/
TS_PREFIX_DIGITS = 5
TS_PREFIX_DIVISOR = 100000  # 10^5; 5-digit prefixes cover ~27h each and cause many list calls
# Hierarchical list: coarse 3-digit prefixes (e.g. 175, 176) cut API calls; optional 2-digit probe (e.g. 17) skips empty bases
TS_PREFIX_COARSE_DIGITS = 3  # Digits used for actual listing; e.g. 175 and 176 cover 17540..17699
TS_PREFIX_PROBE_DIGITS = 2  # Digits for probe prefixes (e.g. 17); MaxKeys=1, skip all 3-digit lists under empty 2-digit prefixes
# Adaptive refine: if a prefix returns a full page (1000) with 0 matches in range, stop paging and recurse with finer prefixes (e.g. 175→1750..1759)
TS_PREFIX_MAX_DIGITS = 10  # Timestamp in filename is 10 digits; stop refining at this length


def load_progress(path: str) -> Optional[dict]:
    """Load checkpoint from progress file; return None if missing or invalid."""
    if not path or not os.path.isfile(path):
        return None
    try:
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    except (json.JSONDecodeError, OSError):
        return None


def save_progress(
    path: str,
    phase: str,
    base_idx: int,
    ts_prefix_idx: int,
    continuation_token: Optional[str],
    total_matched: int,
    total_matched_size: int,
    total_deleted: int,
    total_deleted_size: int,
    last_key: Optional[str],
    no_merge_stats: Optional[dict] = None,
    merged_entries: Optional[list[list[str]]] = None,
) -> None:
    """Write progress file for --resume after crashes. merged_entries is serialized [(cluster_id, component, hour_str), ...]."""
    if not path:
        return
    obj = {
        "phase": phase,
        "base_idx": base_idx,
        "ts_prefix_idx": ts_prefix_idx,
        "continuation_token": continuation_token,
        "total_matched": total_matched,
        "total_matched_size": total_matched_size,
        "total_deleted": total_deleted,
        "total_deleted_size": total_deleted_size,
        "last_key": last_key,
        "updated_at": datetime.now(timezone.utc).isoformat(),
    }
    if no_merge_stats is not None:
        obj["no_merge_stats"] = {
            f"{c}|{p}": {"file_count": s.file_count, "total_size": s.total_size}
            for (c, p), s in no_merge_stats.items()
        }
    if merged_entries is not None:
        # Serialize as [[cluster_id, component, hour_str], ...] for rebuild into a set on resume
        obj["merged_entries"] = [list(t) for t in merged_entries]
    try:
        with open(path, "w", encoding="utf-8") as f:
            json.dump(obj, f, ensure_ascii=False, indent=0)
    except OSError:
        pass


@dataclass
class NoMergeStats:
    """Stats for cluster-component pairs with no matching merged file."""

    file_count: int = 0
    total_size: int = 0

    def add(self, size: int) -> None:
        self.file_count += 1
        self.total_size += size


def unix_to_utc_hour(ts: int) -> str:
    """Convert Unix timestamp to UTC hour string YYYYMMDDHH (same as diagnosis-compute)."""
    return datetime.fromtimestamp(ts, tz=timezone.utc).strftime("%Y%m%d%H")


def get_ts_prefixes(start_ts: int, end_ts: int) -> list[str]:
    """
    Build 5-digit timestamp prefixes for [start_ts, end_ts] (compat / progress display).
    Example: 1771545600..1771631999 -> ["17715", "17716"].
    """
    lo = start_ts // TS_PREFIX_DIVISOR
    hi = end_ts // TS_PREFIX_DIVISOR
    return [str(i) for i in range(lo, hi + 1)]


def get_ts_prefixes_coarse(start_ts: int, end_ts: int, digits: int = TS_PREFIX_COARSE_DIGITS) -> list[str]:
    """
    Build coarse timestamp prefixes (default 3 digits) for [start_ts, end_ts] to cut API calls.
    Example: 1754000000..1764400000 -> 5-digit 17540..17644, deduped 3-digit ["175", "176"];
    each base lists twice (175, 176) instead of 105 times (17540..17644).
    """
    lo = start_ts // TS_PREFIX_DIVISOR
    hi = end_ts // TS_PREFIX_DIVISOR
    divisor = 10 ** (TS_PREFIX_DIGITS - digits)
    seen: set[int] = set()
    out: list[str] = []
    for p in range(lo, hi + 1):
        coarse = p // divisor
        if coarse not in seen:
            seen.add(coarse)
            out.append(str(coarse))
    return out


def parse_small_key(key: str) -> Optional[tuple[str, str, str, str, int]]:
    """
    Parse a small-file key into (tenant, project, cluster_id, component, unix_ts).
    Return None if the format does not match.
    """
    m = SMALL_KEY_PATTERN.match(key)
    if not m:
        return None
    tenant, project, cluster_id, component, ts_str = m.groups()
    try:
        ts = int(ts_str)
    except ValueError:
        return None
    return (tenant, project, cluster_id, component, ts)


def parse_merged_key(key: str) -> Optional[tuple[str, str, str]]:
    """
    Parse a merged-file key into (cluster_id, component, hour_str).
    Supports both path formats.
    """
    m = MERGED_KEY_PATTERN.match(key)
    if m:
        cluster_id, component, hour_str = m.groups()
        return (cluster_id, component, hour_str)
    m = MERGED_KEY_PATTERN_ALT.match(key)
    if m:
        cluster_id, hour_str, component = m.groups()
        return (cluster_id, component, hour_str)
    return None


def get_log_base_prefixes(
    s3_client,
    bucket: str,
    log_progress: bool = True,
    exclude_dirs: Optional[set[str]] = None,
) -> list[str]:
    """
    List all 0/tenant/project/cluster/logs/component/ prefixes under 0/ via Delimiter,
    so later listing can use timestamp prefixes under these bases instead of scanning all of 0/.
    If exclude_dirs is set, skip base paths whose any segment (tenant/project/cluster/component) is in that set.
    """
    result: list[str] = []
    exclude = exclude_dirs or set()

    def list_common_prefixes(prefix: str, depth: int) -> list[str]:
        """Return CommonPrefixes (next-level directories) under this prefix."""
        paginator = s3_client.get_paginator("list_objects_v2")
        prefixes: list[str] = []
        for page in paginator.paginate(
            Bucket=bucket,
            Prefix=prefix,
            Delimiter="/",
            MaxKeys=LIST_PAGE_SIZE,
        ):
            for p in page.get("CommonPrefixes", []):
                prefixes.append(p["Prefix"])
        return prefixes

    # 0/ -> tenants
    tenants = list_common_prefixes("0/", 0)
    if log_progress:
        print(f"  [base] found {len(tenants)} tenants")
    idx_tenant = 0
    for t in tenants:
        idx_tenant += 1
        if exclude and t.rstrip("/").split("/")[-1] in exclude:
            continue
        projects = list_common_prefixes(t, 1)
        for p in projects:
            if exclude and p.rstrip("/").split("/")[-1] in exclude:
                continue
            clusters = list_common_prefixes(p, 2)
            for c in clusters:
                if exclude and c.rstrip("/").split("/")[-1] in exclude:
                    continue
                # 0/t/p/c/logs/
                logs_prefix = c + "logs/"
                components = list_common_prefixes(logs_prefix, 3)
                for comp in components:
                    if exclude and comp.rstrip("/").split("/")[-1] in exclude:
                        continue
                    result.append(comp)  # 0/t/p/c/logs/component/
        if log_progress and idx_tenant % 50 == 0:
            print(f"  [base] scanned tenant {idx_tenant}/{len(tenants)}, {len(result)} base paths so far")
    if log_progress:
        print(f"  [base] done, {len(result)} base paths (0/.../logs/component/)")
    return result


def build_merged_set(
    s3_client,
    bucket: str,
    log_progress: bool = True,
    allowed_hours: Optional[set[str]] = None,
) -> tuple[set[tuple[str, str, str]], int, int]:
    """
    List only objects under diagnosis/data/{clusterid}/merged-logs/, build a set of
    (cluster_id, component, hour_str), and accumulate merged file count/size in the same
    window (for compression ratio vs small files).
    Do not list other paths under diagnosis/data/ (slowlog-stats, log-stats, metrics, etc.).
    If allowed_hours is set, keep only entries whose hour_str is in that set.
    Returns (merged_set, total_merged_count, total_merged_size_bytes).
    """
    merged_set = set()
    total_merged_size = 0
    total_merged_count = 0
    # 1) Use Delimiter to get all diagnosis/data/{clusterid}/
    cluster_prefixes: list[str] = []
    paginator = s3_client.get_paginator("list_objects_v2")
    for page in paginator.paginate(
        Bucket=bucket,
        Prefix="diagnosis/data/",
        Delimiter="/",
        MaxKeys=LIST_PAGE_SIZE,
    ):
        for p in page.get("CommonPrefixes", []):
            cluster_prefixes.append(p["Prefix"])
    if log_progress:
        print(f"  [merged] found {len(cluster_prefixes)} clusters; listing each merged-logs/ only ...")
    # 2) For each cluster, list only the .../merged-logs/ prefix
    total_listed = 0
    for idx, cluster_prefix in enumerate(cluster_prefixes):
        merged_logs_prefix = cluster_prefix + "merged-logs/"
        page_num = 0
        for page in paginator.paginate(
            Bucket=bucket,
            Prefix=merged_logs_prefix,
            MaxKeys=LIST_PAGE_SIZE,
        ):
            page_num += 1
            if "Contents" not in page:
                continue
            n_in_page = len(page["Contents"])
            total_listed += n_in_page
            last_key = page["Contents"][-1]["Key"]
            for obj in page["Contents"]:
                key = obj["Key"]
                size = obj.get("Size", 0) or 0
                parsed = parse_merged_key(key)
                if not parsed:
                    continue
                cluster_id, component, hour_str = parsed
                if allowed_hours is not None and hour_str not in allowed_hours:
                    continue
                merged_set.add(parsed)
                total_merged_count += 1
                total_merged_size += size
            if log_progress and (idx < 3 or page_num <= 2 or n_in_page < LIST_PAGE_SIZE):
                print(
                    f"  [merged] cluster {idx + 1}/{len(cluster_prefixes)} {cluster_prefix.strip('/').split('/')[-1]} "
                    f"page {page_num}: {n_in_page} in page, listed {total_listed} keys total, merged set size {len(merged_set)}"
                )
                print(f"      last key in page: {last_key}")
        if log_progress and (idx + 1) % 100 == 0:
            print(f"  [merged] processed cluster {idx + 1}/{len(cluster_prefixes)}, merged set size {len(merged_set)}")
    return merged_set, total_merged_count, total_merged_size


def _list_one_prefix_page(
    s3_client,
    bucket: str,
    prefix: str,
    continuation_token: Optional[str],
    max_keys: int = LIST_PAGE_SIZE,
) -> dict[str, Any]:
    """One list_objects_v2 call; returns response with Contents/NextContinuationToken."""
    params = {
        "Bucket": bucket,
        "Prefix": prefix,
        "MaxKeys": max_keys,
    }
    if continuation_token:
        params["ContinuationToken"] = continuation_token
    return s3_client.list_objects_v2(**params)


def _probe_prefix_has_keys(
    s3_client,
    bucket: str,
    prefix: str,
) -> bool:
    """Probe whether the prefix has at least one key (MaxKeys=1) for coarse-then-fine listing."""
    page = _list_one_prefix_page(s3_client, bucket, prefix, None, max_keys=1)
    return bool(page.get("Contents"))


def _prefix_timestamp_range(prefix: str) -> tuple[int, int]:
    """
    Given a numeric timestamp prefix (e.g. "175095", 6 digits), return the closed [min_ts, max_ts]
    it covers. Used when refining so only sub-prefixes overlapping [start_ts, end_ts] are listed.
    """
    n = len(prefix)
    if n >= 10:
        t = int(prefix[:10])
        return (t, t)
    base = int(prefix) * (10 ** (10 - n))
    return (base, base + (10 ** (10 - n)) - 1)


def _prefix_overlaps_range(
    prefix: str,
    start_ts: int,
    end_ts: Optional[int],
) -> bool:
    """Return whether the prefix timestamp range overlaps [start_ts, end_ts]."""
    low, high = _prefix_timestamp_range(prefix)
    if end_ts is not None and low > end_ts:
        return False
    if high < start_ts:
        return False
    return True


def list_small_files_with_ts_prefix(
    s3_client,
    bucket: str,
    base_prefixes: list[str],
    ts_prefixes: list[str],
    start_ts: int,
    end_ts: Optional[int],
    on_progress: Optional[
        Callable[
            [
                int,
                int,
                int,
                str,
                int,
                int,
                int,
                dict,
                Optional[str],
                Optional[str],
                Optional[str],
            ],
            None,
        ]
    ] = None,
    progress_info: Optional[dict] = None,
    resume_ctx: Optional[dict] = None,
    probe_2digit: bool = True,
) -> Iterator[tuple[str, int]]:
    """
    List under base_prefixes with ts_prefixes (coarse, e.g. 3-digit 175/176); Prefix = base + ts_prefix.
    Yield only small logs whose filename timestamp is in [start_ts, end_ts].
    If a prefix returns a full page (1000) with 0 matches in range, stop paging and recurse with finer
    prefixes (e.g. 175→1750..1759, 1756→17560..17569).
    If probe_2digit=True, probe each base with a 2-digit prefix (e.g. 17) MaxKeys=1 and skip all
    coarse prefixes under empty 2-digit bases.
    on_progress(..., ts_prefix may be refined, e.g. 1756, 17563).
    resume_ctx: {"base_idx": int, "ts_prefix_idx": int, "continuation_token": str|None} to resume listing.
    """
    total_bases = len(base_prefixes)
    start_base_idx = 0
    start_ts_idx = 0
    start_token: Optional[str] = None
    if resume_ctx:
        start_base_idx = resume_ctx.get("base_idx", 0)
        start_ts_idx = resume_ctx.get("ts_prefix_idx", 0)
        start_token = resume_ctx.get("continuation_token")

    for base_idx in range(start_base_idx, total_bases):
        base = base_prefixes[base_idx]
        probed_empty_2digit: set[str] = set()
        ts_range = range(start_ts_idx, len(ts_prefixes)) if base_idx == start_base_idx else range(len(ts_prefixes))
        for ts_prefix_idx in ts_range:
            ts_prefix_coarse = ts_prefixes[ts_prefix_idx]
            if len(ts_prefix_coarse) >= TS_PREFIX_PROBE_DIGITS:
                two_digit = ts_prefix_coarse[: TS_PREFIX_PROBE_DIGITS]
                if two_digit in probed_empty_2digit:
                    if on_progress:
                        on_progress(
                            base_idx + 1,
                            total_bases,
                            ts_prefix_idx,
                            ts_prefix_coarse,
                            0,
                            0,
                            0,
                            progress_info or {},
                            None,
                            None,
                            None,
                        )
                    continue
                if probe_2digit and (start_token is None or (base_idx, ts_prefix_idx) != (start_base_idx, start_ts_idx)):
                    if not _probe_prefix_has_keys(s3_client, bucket, base + two_digit):
                        probed_empty_2digit.add(two_digit)
                        if on_progress:
                            on_progress(
                                base_idx + 1,
                                total_bases,
                                ts_prefix_idx,
                                ts_prefix_coarse,
                                0,
                                0,
                                0,
                                progress_info or {},
                                None,
                                None,
                                None,
                            )
                        continue

            # Stack: (prefix suffix, continuation_token); empty suffix is the coarse prefix; refine by pushing "0".."9"
            max_suffix_len = TS_PREFIX_MAX_DIGITS - len(ts_prefix_coarse)
            stack: list[tuple[str, Optional[str]]] = [
                ("", start_token if (base_idx == start_base_idx and ts_prefix_idx == start_ts_idx) else None)
            ]
            start_token = None

            page_num = 0
            while stack:
                suffix, continuation_token = stack.pop()
                full_prefix = base + ts_prefix_coarse + suffix
                page = _list_one_prefix_page(s3_client, bucket, full_prefix, continuation_token)
                page_num += 1
                next_token = page.get("NextContinuationToken")
                if "Contents" not in page:
                    if on_progress:
                        display_prefix = ts_prefix_coarse + suffix if suffix else ts_prefix_coarse
                        token_for_progress = next_token if suffix == "" else None
                        on_progress(
                            base_idx + 1,
                            total_bases,
                            ts_prefix_idx,
                            display_prefix,
                            page_num,
                            0,
                            0,
                            progress_info or {},
                            None,
                            None,
                            token_for_progress,
                        )
                    continue
                keys_in_page = len(page["Contents"])
                first_key = page["Contents"][0]["Key"]
                last_key = page["Contents"][-1]["Key"]
                matched_in_page = 0
                for obj in page["Contents"]:
                    key = obj["Key"]
                    size = obj.get("Size", 0)
                    parsed = parse_small_key(key)
                    if not parsed:
                        continue
                    _tenant, _project, _cluster, _component, ts = parsed
                    if ts < start_ts:
                        continue
                    if end_ts is not None and ts > end_ts:
                        continue
                    matched_in_page += 1
                    yield (key, size)
                display_prefix = ts_prefix_coarse + suffix if suffix else ts_prefix_coarse
                # Only top-level (suffix=="") tokens are saved for resume; otherwise resume would be misaligned
                token_for_progress: Optional[str] = next_token if suffix == "" else None
                if on_progress:
                    on_progress(
                        base_idx + 1,
                        total_bases,
                        ts_prefix_idx,
                        display_prefix,
                        page_num,
                        keys_in_page,
                        matched_in_page,
                        progress_info or {},
                        first_key,
                        last_key,
                        token_for_progress,
                    )
                # Full page with 0 matches in range: refine instead of paging; only push sub-prefixes overlapping [start_ts,end_ts]
                if keys_in_page >= LIST_PAGE_SIZE and matched_in_page == 0 and len(suffix) < max_suffix_len:
                    for d in range(9, -1, -1):
                        sub = suffix + str(d)
                        full_digit_prefix = ts_prefix_coarse + sub
                        if _prefix_overlaps_range(full_digit_prefix, start_ts, end_ts):
                            stack.append((sub, None))
                    continue
                if next_token:
                    stack.append((suffix, next_token))


def run(
    bucket: str,
    start_time: datetime,
    end_time: Optional[datetime],
    dry_run: bool,
    region: str,
    skip_delete: bool,
    progress_file: Optional[str] = None,
    resume: bool = False,
    exclude_dirs: Optional[set[str]] = None,
) -> None:
    if boto3 is None:
        print("Please install boto3: pip install boto3", file=sys.stderr)
        sys.exit(1)

    start_ts = int(start_time.timestamp())
    end_ts_int = int(end_time.timestamp()) if end_time else None
    print(f"Time range: start_ts={start_ts} ({start_time.isoformat()})", end="")
    if end_ts_int is not None:
        print(f", end_ts={end_ts_int} ({end_time.isoformat()})")
    else:
        print(", end_ts=unbounded")
    if end_ts_int is None:
        print(
            "Error: --end-time is required for timestamp-prefix listing; without it a full scan of 0/ would use too many API calls.",
            file=sys.stderr,
        )
        sys.exit(1)
    s3_client = boto3.client("s3", region_name=region)

    resume_ctx: Optional[dict] = None
    total_deleted_count = 0
    total_deleted_size = 0
    no_merge_stats: dict[tuple[str, str], NoMergeStats] = defaultdict(NoMergeStats)
    if progress_file and resume:
        checkpoint = load_progress(progress_file)
        if checkpoint and checkpoint.get("phase") == "list":
            resume_ctx = {
                "base_idx": checkpoint.get("base_idx", 0),
                "ts_prefix_idx": checkpoint.get("ts_prefix_idx", 0),
                "continuation_token": checkpoint.get("continuation_token"),
            }
            total_deleted_count = checkpoint.get("total_deleted", 0)
            total_deleted_size = checkpoint.get("total_deleted_size", 0)
            # Restore no_merge_stats (resume continues accumulating in the loop)
            for k, v in (checkpoint.get("no_merge_stats") or {}).items():
                parts = k.split("|", 1)
                if len(parts) == 2:
                    no_merge_stats[(parts[0], parts[1])].file_count = v.get("file_count", 0)
                    no_merge_stats[(parts[0], parts[1])].total_size = v.get("total_size", 0)
            print(f"Resumed from progress file: {progress_file}, continuing at base_idx={resume_ctx['base_idx']} ts_prefix_idx={resume_ctx['ts_prefix_idx']}")
    if resume_ctx and progress_file:
        ck = load_progress(progress_file)
        total_small_count_init = ck.get("total_matched", 0) if ck else 0
        total_small_size_init = ck.get("total_matched_size", 0) if ck else 0
    else:
        total_small_count_init = 0
        total_small_size_init = 0

    # Keep only hours in the time range to shrink the merged set
    allowed_hours = set()
    t = start_ts
    while t < end_ts_int:
        allowed_hours.add(unix_to_utc_hour(t))
        t += 3600
    print(f"{len(allowed_hours)} hours in time range (YYYYMMDDHH)")

    # 1) Build merged set (on resume, load from progress file instead of listing diagnosis/data/)
    merged_set: set[tuple[str, str, str]]
    total_merged_count: Optional[int] = None
    total_merged_size: Optional[int] = None
    if progress_file and resume:
        ck = load_progress(progress_file)
        if ck and ck.get("merged_entries"):
            merged_set = set(tuple(e) for e in ck["merged_entries"] if len(e) == 3)
            print(f"Phase 1/3: loaded merged set from progress file, {len(merged_set)} (cluster, component, hour); skip list")
        else:
            print("Phase 1/3: listing merged files (diagnosis/data/...) ...")
            merged_set, total_merged_count, total_merged_size = build_merged_set(
                s3_client, bucket, allowed_hours=allowed_hours
            )
            print(f"  [merged] done, {len(merged_set)} (cluster, component, hour) have merged data in range")
    else:
        print("Phase 1/3: listing merged files (diagnosis/data/...) ...")
        merged_set, total_merged_count, total_merged_size = build_merged_set(
            s3_client, bucket, allowed_hours=allowed_hours
        )
        print(f"  [merged] done, {len(merged_set)} (cluster, component, hour) have merged data in range")

    # 2) Discover all 0/.../logs/component/ base paths (optional excludes)
    print("Phase 2/3: discovering small-log base paths (0/.../logs/component/) ...")
    base_prefixes = get_log_base_prefixes(
        s3_client, bucket, exclude_dirs=exclude_dirs
    )
    if not base_prefixes:
        print("  No base paths found; exiting.")
        return
    ts_prefixes = get_ts_prefixes_coarse(start_ts, end_ts_int)
    print(
        f"  Will list by coarse timestamp prefixes: {ts_prefixes} ({len(ts_prefixes)} prefixes, {TS_PREFIX_COARSE_DIGITS} digits); "
        f"probe each base with {TS_PREFIX_PROBE_DIGITS} digits and skip if empty"
    )

    # 3) List small files by base + ts_prefix; stream deletes (flush every 1000) and write progress
    to_delete_batch: list[tuple[str, int]] = []  # flush at most 1000 keys
    sample_to_delete: list[str] = []  # dry_run: keep up to 20 keys for printing
    total_small_count = total_small_count_init
    total_small_size = total_small_size_init
    total_list_pages = 0
    progress_info = {"count": total_small_count, "size": total_small_size}
    last_saved_base_idx = -1
    last_saved_ts_prefix_idx = -1
    last_saved_token: Optional[str] = None

    def _on_progress(
        base_idx_1: int,
        total_bases: int,
        ts_prefix_idx: int,
        ts_prefix: str,
        page_num: int,
        keys_in_page: int,
        matched_in_page: int,
        info: dict,
        first_key: Optional[str],
        last_key: Optional[str],
        next_continuation_token: Optional[str],
    ) -> None:
        nonlocal total_list_pages, last_saved_base_idx, last_saved_ts_prefix_idx, last_saved_token
        total_list_pages += 1
        short_base = (
            base_prefixes[base_idx_1 - 1].rstrip("/").split("/")[-1]
            if base_idx_1 <= len(base_prefixes)
            else "?"
        )
        cnt, sz = info.get("count", 0), info.get("size", 0)
        print(
            f"  [small] base {base_idx_1}/{total_bases} ({short_base}) ts_prefix={ts_prefix} "
            f"page {page_num}: {keys_in_page} keys in page, {matched_in_page} in range | matched so far: {cnt}, {sz / (1024*1024):.1f} MB"
        )
        if last_key is not None:
            print(f"      last key in page: {last_key}")
        last_saved_base_idx = base_idx_1 - 1
        last_saved_ts_prefix_idx = ts_prefix_idx
        last_saved_token = next_continuation_token
        if progress_file:
            save_progress(
                progress_file,
                phase="list",
                base_idx=last_saved_base_idx,
                ts_prefix_idx=last_saved_ts_prefix_idx,
                continuation_token=last_saved_token,
                total_matched=cnt,
                total_matched_size=sz,
                total_deleted=total_deleted_count,
                total_deleted_size=total_deleted_size,
                last_key=last_key,
                no_merge_stats=no_merge_stats,
                merged_entries=[list(x) for x in merged_set],
            )

    print("Phase 3/3: listing small files by base path + timestamp prefix (streaming delete, progress on disk)...")
    for key, size in list_small_files_with_ts_prefix(
        s3_client,
        bucket,
        base_prefixes,
        ts_prefixes,
        start_ts,
        end_ts_int,
        on_progress=_on_progress,
        progress_info=progress_info,
        resume_ctx=resume_ctx,
    ):
        total_small_count += 1
        total_small_size += size
        progress_info["count"] = total_small_count
        progress_info["size"] = total_small_size
        parsed = parse_small_key(key)
        if not parsed:
            continue
        _tenant, _project, cluster_id, component, ts = parsed
        hour_str = unix_to_utc_hour(ts)
        if (cluster_id, component, hour_str) in merged_set:
            if skip_delete or dry_run:
                total_deleted_count += 1
                total_deleted_size += size
                if dry_run and len(sample_to_delete) < 20:
                    sample_to_delete.append(key)
            else:
                to_delete_batch.append((key, size))
                if len(to_delete_batch) >= DELETE_BATCH_SIZE:
                    keys_batch = [k for k, _ in to_delete_batch]
                    sizes_batch = sum(s for _, s in to_delete_batch)
                    s3_client.delete_objects(
                        Bucket=bucket,
                        Delete={"Objects": [{"Key": k} for k in keys_batch], "Quiet": True},
                    )
                    total_deleted_count += len(keys_batch)
                    total_deleted_size += sizes_batch
                    to_delete_batch.clear()
                    print(f"  [delete] streamed delete of {len(keys_batch)} keys; deleted so far: {total_deleted_count}")
        else:
            no_merge_stats[(cluster_id, component)].add(size)
    # flush remaining deletes
    if to_delete_batch and not skip_delete and not dry_run:
        keys_batch = [k for k, _ in to_delete_batch]
        sizes_batch = sum(s for _, s in to_delete_batch)
        s3_client.delete_objects(
            Bucket=bucket,
            Delete={"Objects": [{"Key": k} for k in keys_batch], "Quiet": True},
        )
        total_deleted_count += len(keys_batch)
        total_deleted_size += sizes_batch
        print(f"  [delete] deleted final batch of {len(keys_batch)} keys; deleted so far: {total_deleted_count}")
    if progress_file:
        save_progress(
            progress_file,
            phase="done",
            base_idx=last_saved_base_idx,
            ts_prefix_idx=last_saved_ts_prefix_idx,
            continuation_token=None,
            total_matched=total_small_count,
            total_matched_size=total_small_size,
            total_deleted=total_deleted_count,
            total_deleted_size=total_deleted_size,
            last_key=None,
            no_merge_stats=no_merge_stats,
            merged_entries=[list(x) for x in merged_set],
        )
    print(f"  [small] done; ~{total_list_pages} list pages, matched {total_small_count} small files")

    # 4) Print stats
    deleted_count = total_deleted_count
    deleted_size = total_deleted_size
    print()
    print("======== Small-file overall stats (for post-merge compression estimate) ========")
    print(f"  Small file count: {total_small_count}")
    print(f"  Small file total size (bytes): {total_small_size}")
    print(f"  Small file total size (MB):    {total_small_size / (1024*1024):.2f}")
    if total_merged_size is not None and total_merged_count is not None:
        print()
        print("======== Merged-file stats in the same time window ========")
        print(f"  Merged file count: {total_merged_count}")
        print(f"  Merged total size (bytes): {total_merged_size}")
        print(f"  Merged total size (MB):    {total_merged_size / (1024*1024):.2f}")
        if total_merged_size > 0 and deleted_size > 0:
            ratio = deleted_size / total_merged_size
            print(f"  Deletable-small / merged size ratio: {ratio:.2f}x (merged is ~{total_merged_size/deleted_size*100:.1f}% of deletable small files)")
    else:
        print("  (Compare with merged size in the same window for compression gain; no merged stats if merged listing was skipped)")
    print()
    print("======== Deletable (merged already exists) ========")
    print(f"  Deletable file count: {deleted_count}")
    print(f"  Deletable total size (bytes): {deleted_size}")
    print(f"  Deletable total size (MB):    {deleted_size / (1024*1024):.2f}")
    print()
    print("======== No matching merged (by cluster-component) ========")
    if no_merge_stats:
        for (cluster_id, component), st in sorted(no_merge_stats.items()):
            print(f"  cluster={cluster_id} component={component}  file_count={st.file_count}  total_size_bytes={st.total_size}  total_size_MB={st.total_size/(1024*1024):.2f}")
    else:
        print("  (none)")
    print()

    # 5) In dry-run, print sample keys that would be deleted (real deletes already streamed in phase 3)
    if dry_run and total_deleted_count > 0:
        print("[dry-run] Sample keys that would be deleted (DeleteObjects not called):")
        for k in sample_to_delete[:20]:
            print(f"  {k}")
        if total_deleted_count > 20:
            print(f"  ... {total_deleted_count} total")
    elif skip_delete and total_deleted_count > 0:
        print("Delete skipped (--skip-delete); deletable counts above are stats only.")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Clean up small S3 logs that have already been merged, and report small-file count/size."
    )
    parser.add_argument(
        "--bucket",
        required=True,
        help="S3 bucket name, e.g. o11y-prod-shared-us-east-1",
    )
    parser.add_argument(
        "--start-time",
        required=True,
        help="Only process small files with timestamp >= this time. Format: 2026-02-20 or 2026-02-20T00:00:00Z",
    )
    parser.add_argument(
        "--end-time",
        default=None,
        help="Required. Only process files with timestamp <= this time; also used for timestamp-prefix listing. Without it, a full scan of 0/ would explode API usage. Same format as --start-time.",
    )
    parser.add_argument(
        "--region",
        default="us-east-1",
        help="S3 region (default us-east-1)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Only report and list keys that would be deleted; do not call DeleteObjects",
    )
    parser.add_argument(
        "--skip-delete",
        action="store_true",
        help="List and stats only; do not delete (unlike dry-run: does not print keys to delete)",
    )
    parser.add_argument(
        "--progress-file",
        default=None,
        metavar="PATH",
        help="Progress file path. Written after each list page; use --resume to continue after a crash without full rescan.",
    )
    parser.add_argument(
        "--resume",
        action="store_true",
        help="Resume from --progress-file (must be used together with --progress-file).",
    )
    parser.add_argument(
        "--exclude-dirs",
        nargs="+",
        default=None,
        metavar="NAME",
        help="Directory names to skip in phase 2 (any matching path segment). Example: --exclude-dirs prod-devtier dev-devtier 30008 30073",
    )
    args = parser.parse_args()

    def parse_dt(s: str) -> datetime:
        s = s.strip()
        for fmt in ("%Y-%m-%dT%H:%M:%SZ", "%Y-%m-%d %H:%M:%S", "%Y-%m-%d"):
            try:
                return datetime.strptime(s, fmt).replace(tzinfo=timezone.utc)
            except ValueError:
                continue
        raise ValueError(f"Unable to parse datetime: {s}")

    start_time = parse_dt(args.start_time)
    end_time = parse_dt(args.end_time) if args.end_time else None
    if end_time is not None and end_time < start_time:
        print("--end-time must be >= --start-time", file=sys.stderr)
        sys.exit(1)

    if args.resume and not args.progress_file:
        print("--resume requires --progress-file.", file=sys.stderr)
        sys.exit(1)
    exclude_dirs = set(args.exclude_dirs) if args.exclude_dirs else None
    if exclude_dirs:
        print(f"Excluded dirs (skipped in phase 2): {sorted(exclude_dirs)}")
    run(
        bucket=args.bucket,
        start_time=start_time,
        end_time=end_time,
        dry_run=args.dry_run,
        region=args.region,
        skip_delete=args.skip_delete,
        progress_file=args.progress_file,
        resume=args.resume,
        exclude_dirs=exclude_dirs,
    )


if __name__ == "__main__":
    main()