"""
This script connects to a NATS server, deletes a specified consumer from the 'tasks' stream,
and purges all messages from the stream. It is intended to be used as a cleanup utility
to remove stale consumers and reset the queue to a clean state.

Usage:
    python scripts/drain_nats.py --consumer <consumer_name> [--purge]
"""

import asyncio
import logging
import argparse

import nats
from nats.aio.client import Client as NATSClient

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

# Attempt to import settings
try:
    from preloop.config import settings
except ImportError:
    logging.error(
        "Could not import settings from preloop.config. "
        "Please ensure that the script is run from the root of the project "
        "and the project is installed in editable mode (e.g., 'pip install -e .')."
    )
    exit(1)


async def drain_queue(consumer_name: str, purge: bool):
    """Connects to NATS, deletes the consumer, and optionally purges the stream."""
    nc: NATSClient | None = None
    try:
        logging.info(f"Connecting to NATS at {settings.nats_url}...")
        nc = await nats.connect(settings.nats_url, name="nats-drainer")
        js = nc.jetstream()
        logging.info("Connection successful.")

        # Delete the specified consumer
        try:
            await js.delete_consumer(stream="tasks", consumer=consumer_name)
            logging.info(f"Successfully deleted consumer '{consumer_name}'.")
        except Exception as e:
            logging.error(
                f"Could not delete consumer '{consumer_name}'. It may not exist. Error: {e}"
            )

        # Purge the stream if requested
        if purge:
            try:
                await js.purge_stream("tasks")
                logging.info(
                    "Successfully purged all messages from the 'tasks' stream."
                )
            except Exception as e:
                logging.error(f"Could not purge the 'tasks' stream. Error: {e}")

    except nats.errors.NoServersError:
        logging.error(f"Could not connect to NATS at {settings.nats_url}. Aborting.")
    except Exception as e:
        logging.error(f"An error occurred during the drain process: {e}")
    finally:
        if nc and not nc.is_closed:
            logging.info("Closing NATS connection...")
            await nc.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="NATS stream cleanup utility.")
    parser.add_argument(
        "--consumer",
        type=str,
        required=True,
        help="The name of the consumer to delete.",
    )
    parser.add_argument(
        "--purge",
        action="store_true",
        help="If set, purge all messages from the 'tasks' stream after deleting the consumer.",
    )
    args = parser.parse_args()

    try:
        asyncio.run(drain_queue(args.consumer, args.purge))
    except KeyboardInterrupt:
        logging.info("\nDrain process interrupted by user. Exiting.")
