#!/usr/bin/env python
"""
Script to start Celery workers and beat scheduler
Usage:
    python scripts/start_workers.py worker     # Start worker only
    python scripts/start_workers.py beat       # Start beat only
    python scripts/start_workers.py all        # Start both
    python scripts/start_workers.py flower     # Start Flower monitoring
"""
import os
import sys
import subprocess
import signal
from typing import List

# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


def start_worker(queues: List[str] = None, concurrency: int = 4):
    """Start Celery worker"""
    cmd = [
        "celery",
        "-A", "celery_worker",
        "worker",
        f"--concurrency={concurrency}",
        "--loglevel=info",
        "--pool=prefork",
    ]

    if queues:
        cmd.extend(["-Q", ",".join(queues)])

    print(f"Starting worker: {' '.join(cmd)}")
    return subprocess.Popen(cmd)


def start_beat():
    """Start Celery beat scheduler"""
    cmd = [
        "celery",
        "-A", "celery_worker",
        "beat",
        "--loglevel=info",
    ]

    print(f"Starting beat: {' '.join(cmd)}")
    return subprocess.Popen(cmd)


def start_flower(port: int = 5555):
    """Start Flower monitoring UI"""
    cmd = [
        "celery",
        "-A", "celery_worker",
        "flower",
        f"--port={port}",
        "--broker=redis://localhost:6379/0",
    ]

    print(f"Starting Flower: {' '.join(cmd)}")
    return subprocess.Popen(cmd)


def main():
    processes = []

    if len(sys.argv) < 2:
        print("Usage: python start_workers.py [worker|beat|all|flower]")
        sys.exit(1)

    mode = sys.argv[1].lower()

    try:
        if mode == "worker":
            # Start multiple workers for different queues
            processes.append(start_worker(
                queues=["default", "high_priority"],
                concurrency=2
            ))
            processes.append(start_worker(
                queues=["emails"],
                concurrency=2
            ))
            processes.append(start_worker(
                queues=["invoices", "analytics"],
                concurrency=2
            ))
            processes.append(start_worker(
                queues=["low_priority"],
                concurrency=1
            ))

        elif mode == "beat":
            processes.append(start_beat())

        elif mode == "all":
            # Start all workers and beat
            processes.append(start_worker(
                queues=["default", "high_priority"],
                concurrency=2
            ))
            processes.append(start_worker(
                queues=["emails"],
                concurrency=2
            ))
            processes.append(start_worker(
                queues=["invoices", "analytics"],
                concurrency=2
            ))
            processes.append(start_worker(
                queues=["low_priority"],
                concurrency=1
            ))
            processes.append(start_beat())

        elif mode == "flower":
            processes.append(start_flower())

        else:
            print(f"Unknown mode: {mode}")
            sys.exit(1)

        # Wait for processes
        def signal_handler(sig, frame):
            print("\nShutting down workers...")
            for p in processes:
                p.terminate()
            sys.exit(0)

        signal.signal(signal.SIGINT, signal_handler)
        signal.signal(signal.SIGTERM, signal_handler)

        print(f"\n{len(processes)} process(es) started. Press Ctrl+C to stop.")

        for p in processes:
            p.wait()

    except KeyboardInterrupt:
        print("\nShutting down...")
        for p in processes:
            p.terminate()


if __name__ == "__main__":
    main()