#!/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
SERVICES_DIR = "/mnt/storage-box/services"

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 cmd_start(args):
    """Start a Docker service"""
    return run_docker_compose(["up", "-d"], service_name=args.service)

def cmd_stop(args):
    """Stop a Docker service"""
    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 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 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:")
    for service in sorted(services):
        println(f"  - {service}")
    
    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_parser.add_argument("service", help="Service to start")
    
    # Stop command
    stop_parser = subparsers.add_parser("stop", help="Stop a Docker service")
    stop_parser.add_argument("service", help="Service to stop")
    
    # Restart command
    restart_parser = subparsers.add_parser("restart", help="Restart a Docker service")
    restart_parser.add_argument("service", help="Service to restart")
    
    # 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
    subparsers.add_parser("list", help="List available Docker services")
    
    # 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,
        "ps": cmd_ps,
        "logs": cmd_logs,
        "list": cmd_list
    }
    
    return commands[args.command](args)

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