#!/usr/bin/env python3

import os
import sys
import subprocess
import argparse

# Import helper functions
sys.path.append(os.path.join(os.path.expanduser("~/.dotfiles"), "bin"))
from helpers.functions import printfe, println, logo

# Base directory for Docker services $HOME/services
SERVICES_DIR = os.path.join(os.path.expanduser("~"), "services")
# Protected services that should never be stopped
PROTECTED_SERVICES = ["juicefs-redis"]


def get_service_path(service_name):
    """Return the path to a service's docker-compose file"""
    service_dir = os.path.join(SERVICES_DIR, service_name)
    compose_file = os.path.join(service_dir, "docker-compose.yml")

    if not os.path.exists(compose_file):
        printfe("red", f"Error: Service '{service_name}' not found at {compose_file}")
        return None

    return compose_file


def run_docker_compose(args, service_name=None, compose_file=None):
    """Run docker compose command with provided args"""
    if service_name and not compose_file:
        compose_file = get_service_path(service_name)
        if not compose_file:
            return 1

    cmd = ["docker", "compose"]

    if compose_file:
        cmd.extend(["-f", compose_file])

    cmd.extend(args)

    printfe("blue", f"Running: {' '.join(cmd)}")
    result = subprocess.run(cmd)
    return result.returncode


def get_all_services():
    """Return a list of all available services"""
    if not os.path.exists(SERVICES_DIR):
        return []

    services = [
        d
        for d in os.listdir(SERVICES_DIR)
        if os.path.isdir(os.path.join(SERVICES_DIR, d))
        and os.path.exists(os.path.join(SERVICES_DIR, d, "docker-compose.yml"))
    ]

    return sorted(services)


def cmd_start(args):
    """Start a Docker service"""
    if args.all:
        services = get_all_services()
        if not services:
            printfe("yellow", "No services found to start")
            return 0

        printfe("blue", f"Starting all services: {', '.join(services)}")

        failed_services = []
        for service in services:
            printfe("blue", f"\n=== Starting {service} ===")
            result = run_docker_compose(["up", "-d"], service_name=service)
            if result != 0:
                failed_services.append(service)

        if failed_services:
            printfe(
                "red",
                f"\nFailed to start the following services: {', '.join(failed_services)}",
            )
            return 1
        else:
            printfe("green", "\nAll services started successfully")
            return 0
    else:
        return run_docker_compose(["up", "-d"], service_name=args.service)


def cmd_stop(args):
    """Stop a Docker service"""
    if args.all:
        running_services = get_all_running_services()
        if not running_services:
            printfe("yellow", "No running services found to stop")
            return 0

        # Filter out the protected services
        safe_services = [s for s in running_services if s not in PROTECTED_SERVICES]

        # Check if protected services were filtered out
        protected_running = [s for s in running_services if s in PROTECTED_SERVICES]
        if protected_running:
            printfe(
                "yellow",
                f"Note: {', '.join(protected_running)} will not be stopped as they are protected services",
            )

        if not safe_services:
            printfe(
                "yellow", "No services to stop (all running services are protected)"
            )
            return 0

        printfe("blue", f"Stopping all running services: {', '.join(safe_services)}")

        failed_services = []
        for service in safe_services:
            printfe("blue", f"\n=== Stopping {service} ===")
            result = run_docker_compose(["down"], service_name=service)
            if result != 0:
                failed_services.append(service)

        if failed_services:
            printfe(
                "red",
                f"\nFailed to stop the following services: {', '.join(failed_services)}",
            )
            return 1
        else:
            printfe("green", "\nAll running services stopped successfully")
            return 0
    else:
        # Check if trying to stop a protected service
        if args.service in PROTECTED_SERVICES:
            printfe(
                "red",
                f"Error: {args.service} is a protected service and cannot be stopped",
            )
            printfe(
                "yellow",
                f"The {args.service} service is required for other services to work properly",
            )
            return 1
        return run_docker_compose(["down"], service_name=args.service)


def cmd_restart(args):
    """Restart a Docker service"""
    return run_docker_compose(["restart"], service_name=args.service)


def get_all_running_services():
    """Return a list of all running services"""
    if not os.path.exists(SERVICES_DIR):
        return []

    running_services = []
    services = [
        d
        for d in os.listdir(SERVICES_DIR)
        if os.path.isdir(os.path.join(SERVICES_DIR, d))
        and os.path.exists(os.path.join(SERVICES_DIR, d, "docker-compose.yml"))
    ]

    for service in services:
        if check_service_running(service) > 0:
            running_services.append(service)

    return running_services


def cmd_update(args):
    """Update a Docker service by pulling new images and recreating containers if needed"""
    if args.all:
        running_services = get_all_running_services()
        if not running_services:
            printfe("yellow", "No running services found to update")
            return 0

        printfe("blue", f"Updating all running services: {', '.join(running_services)}")

        failed_services = []
        for service in running_services:
            printfe("blue", f"\n=== Updating {service} ===")

            # Pull the latest images
            pull_result = run_docker_compose(["pull"], service_name=service)

            # Bring the service up with the latest images
            up_result = run_docker_compose(["up", "-d"], service_name=service)

            if pull_result != 0 or up_result != 0:
                failed_services.append(service)

        if failed_services:
            printfe(
                "red",
                f"\nFailed to update the following services: {', '.join(failed_services)}",
            )
            return 1
        else:
            printfe("green", "\nAll running services updated successfully")
            return 0
    else:
        # The original single-service update logic
        # First pull the latest images
        pull_result = run_docker_compose(["pull"], service_name=args.service)
        if pull_result != 0:
            return pull_result

        # Then bring the service up with the latest images
        return run_docker_compose(["up", "-d"], service_name=args.service)


def cmd_ps(args):
    """Show Docker service status"""
    if args.service:
        return run_docker_compose(["ps"], service_name=args.service)
    else:
        return run_docker_compose(["ps"])


def cmd_logs(args):
    """Show Docker service logs"""
    cmd = ["logs"]

    if args.follow:
        cmd.append("-f")

    if args.tail:
        cmd.extend(["--tail", args.tail])

    return run_docker_compose(cmd, service_name=args.service)


def check_service_running(service_name):
    """Check if service has running containers and return the count"""
    compose_file = get_service_path(service_name)
    if not compose_file:
        return 0

    result = subprocess.run(
        ["docker", "compose", "-f", compose_file, "ps", "--quiet"],
        capture_output=True,
        text=True,
    )

    # Count non-empty lines to get container count
    containers = [line for line in result.stdout.strip().split("\n") if line]
    return len(containers)


def cmd_list(args):
    """List available Docker services"""
    if not os.path.exists(SERVICES_DIR):
        printfe("red", f"Error: Services directory not found at {SERVICES_DIR}")
        return 1

    services = [
        d
        for d in os.listdir(SERVICES_DIR)
        if os.path.isdir(os.path.join(SERVICES_DIR, d))
        and os.path.exists(os.path.join(SERVICES_DIR, d, "docker-compose.yml"))
    ]

    if not services:
        printfe("yellow", "No Docker services found")
        return 0

    println("Available Docker services:", "blue")
    for service in sorted(services):
        container_count = check_service_running(service)
        is_running = container_count > 0

        if is_running:
            status = f"[RUNNING - {container_count} container{'s' if container_count > 1 else ''}]"
            color = "green"
        else:
            status = "[STOPPED]"
            color = "red"

        printfe(color, f"  - {service:<20} {status}")

    return 0


def main():
    parser = argparse.ArgumentParser(description="Manage Docker services")
    subparsers = parser.add_subparsers(dest="command", help="Command to run")

    # Start command
    start_parser = subparsers.add_parser("start", help="Start a Docker service")
    start_group = start_parser.add_mutually_exclusive_group(required=True)
    start_group.add_argument("--all", action="store_true", help="Start all services")
    start_group.add_argument("service", nargs="?", help="Service to start")
    start_group.add_argument(
        "--service", dest="service", help="Service to start (deprecated)"
    )

    # Stop command
    stop_parser = subparsers.add_parser("stop", help="Stop a Docker service")
    stop_group = stop_parser.add_mutually_exclusive_group(required=True)
    stop_group.add_argument(
        "--all", action="store_true", help="Stop all running services"
    )
    stop_group.add_argument("service", nargs="?", help="Service to stop")
    stop_group.add_argument(
        "--service", dest="service", help="Service to stop (deprecated)"
    )

    # Restart command
    restart_parser = subparsers.add_parser("restart", help="Restart a Docker service")
    restart_parser.add_argument("service", help="Service to restart")

    # Update command
    update_parser = subparsers.add_parser(
        "update",
        help="Update a Docker service (pull new images and recreate if needed)",
    )
    update_group = update_parser.add_mutually_exclusive_group(required=True)
    update_group.add_argument(
        "--all", action="store_true", help="Update all running services"
    )
    update_group.add_argument("service", nargs="?", help="Service to update")
    update_group.add_argument(
        "--service", dest="service", help="Service to update (deprecated)"
    )

    # PS command
    ps_parser = subparsers.add_parser("ps", help="Show Docker service status")
    ps_parser.add_argument("service", nargs="?", help="Service to check")

    # Logs command
    logs_parser = subparsers.add_parser("logs", help="Show Docker service logs")
    logs_parser.add_argument("service", help="Service to show logs for")
    logs_parser.add_argument(
        "-f", "--follow", action="store_true", help="Follow log output"
    )
    logs_parser.add_argument(
        "--tail", help="Number of lines to show from the end of logs"
    )

    # List command and its alias
    subparsers.add_parser("list", help="List available Docker services")
    subparsers.add_parser("ls", help="List available Docker services (alias for list)")

    # Parse arguments
    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        return 1

    # Execute the appropriate command
    commands = {
        "start": cmd_start,
        "stop": cmd_stop,
        "restart": cmd_restart,
        "update": cmd_update,
        "ps": cmd_ps,
        "logs": cmd_logs,
        "list": cmd_list,
        "ls": cmd_list,  # Alias 'ls' to the same function as 'list'
    }

    return commands[args.command](args)


if __name__ == "__main__":
    sys.exit(main())