#!/home/thebears/envs/vector_search/bin/python

"""Trigger and (optionally) wait on Milvus compaction for one or more collections.

Each collection is flushed before compaction is triggered, since Milvus only
compacts sealed segments -- flushing seals any growing segments so recently
inserted data is actually eligible.

By default this compacts every collection on the server once and exits, which is
the right shape for a cron job / systemd oneshot timer. Pass --loop to instead run
forever, re-compacting every --interval seconds (for a long-running systemd service).
"""

from pymilvus import MilvusClient
from fnmatch import fnmatch
import argparse
import logging
import time
import random

logger = logging.getLogger(__name__)

DEFAULT_URI = "http://192.168.1.242:19530"


def compact_collections(
    client,
    pattern="*",
    is_clustering=False,
    wait=True,
    poll_interval=60,
    timeout=3600*10,
    dry_run=False,
):
    """Trigger compaction for every collection whose name matches `pattern`.

    Returns a dict of {collection_name: True/False/None}, where None means the
    job was triggered but not waited on (wait=False).
    """
    results = dict()
    collections = [c for c in client.list_collections() if fnmatch(c, pattern)]

    if not collections:
        logger.info(f"No collections matched pattern {pattern!r}")
        return results

    random.shuffle(collections)
    for coll_name in collections:
        if dry_run:
            logger.info(f"[dry-run] would flush and compact {coll_name}")
            continue

        try:
            client.flush(coll_name)
        except Exception as e:
            logger.error(f"Failed to flush {coll_name} before compaction: {e}")
            results[coll_name] = False
            continue

        try:
            job_id = client.compact(coll_name, is_clustering=is_clustering)
        except Exception as e:
            logger.error(f"Failed to trigger compaction for {coll_name}: {e}")
            results[coll_name] = False
            continue

        logger.info(f"Triggered compaction job {job_id} for {coll_name}")

        if not wait:
            results[coll_name] = None
            continue

        waited = 0
        state = None
        while waited < timeout:
            try:
                state = client.get_compaction_state(job_id)
            except Exception as e:
                logger.error(f"Failed to poll compaction state for {coll_name} (job {job_id}): {e}")
                break
            if state == "Completed":
                break
            time.sleep(poll_interval)
            waited += poll_interval

        success = state == "Completed"
        if success:
            logger.info(f"Compaction of {coll_name} completed (job {job_id})")
        else:
            logger.warning(f"Compaction of {coll_name} (job {job_id}) ended with state={state} after {waited}s")
        results[coll_name] = success

    return results


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog="Auto-compact Milvus collections"
    )
    parser.add_argument("--uri", default=DEFAULT_URI, help="Milvus server URI")
    parser.add_argument("--pattern", default="*", help="fnmatch glob to select which collections to compact")
    parser.add_argument("--is-clustering", action="store_true", dest="is_clustering",
                         help="Trigger clustering compaction instead of a regular merge compaction")
    parser.add_argument("--no-wait", action="store_false", dest="wait",
                         help="Trigger compaction jobs without waiting for them to finish")
    parser.add_argument("--poll-interval", type=float, default=5, help="Seconds between compaction state polls")
    parser.add_argument("--timeout", type=float, default=3600, help="Max seconds to wait per collection's compaction job")
    parser.add_argument("--dry-run", action="store_true", help="List collections that would be compacted and exit")
    parser.add_argument("--loop", action="store_true", help="Run forever, re-compacting on a fixed interval")
    parser.add_argument("--interval", type=float, default=6 * 3600, help="Seconds between runs when --loop is set")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(asctime)s|%(levelname)s|%(message)s")

    client = MilvusClient(uri=args.uri)
    try:
        while True:
            results = compact_collections(
                client,
                pattern=args.pattern,
                is_clustering=args.is_clustering,
                wait=args.wait,
                poll_interval=args.poll_interval,
                timeout=args.timeout,
                dry_run=args.dry_run,
            )
            for coll_name, ok in results.items():
                if ok is False:
                    print(f"Failed: {coll_name}")

            if not args.loop:
                break
            logger.info(f"Sleeping {args.interval}s before next compaction pass")
            time.sleep(args.interval)
    finally:
        client.close()
