import argparse
import datetime
import logging
import os
import subprocess
from pathlib import Path
import docker
from pytz import timezone

# Set up logging
LOG_FILE = Path("/mnt/docker-persistent-data/docker/update-containers.log")
LOG_FILE.parent.mkdir(parents=True, exist_ok=True)

# Clear the log file by opening it in write mode before configuring logging
with open(LOG_FILE, 'w'):
    pass

logging.basicConfig(filename=LOG_FILE, level=logging.INFO, format='%(asctime)s %(levelname)s: %(message)s', datefmt='%H:%M:%S')

DEFAULT_STACKS = ["traefik", "nextcloud", "passbolt", "searxng", "gitea", "gotify", "grafana", "gramps", "portainer", "prometheus", "uptime-kuma"]
DOCKER_PATH = Path.home() / "docker"

def run_compose_command(service_dir: Path, command: list[str]):
    """Run a docker compose CLI command in the given directory."""
    try:
        result = subprocess.run(
            ["docker", "compose"] + command,
            cwd=service_dir,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            check=True
        )
        logging.info(result.stdout)
    except subprocess.CalledProcessError as e:
        logging.error(f"Error running {' '.join(e.cmd)}:\n{e.output}")
        raise

def check_container_status(docker_client, service_name):
    containers = docker_client.containers.list(all=True, filters={'name': service_name})
    for container in containers:
        container_info = docker_client.api.inspect_container(container.id)
        container_state = container_info['State']
        if not container_state['Running']:
            logging.error(f"Container {container.name} for service {service_name} is not running. Status: {container_state['Status']}")

def main(stacks):
    current_time = datetime.datetime.now(timezone("Australia/Brisbane"))
    logging.info("=== Update started: %s ===", current_time.strftime("%Y-%m-%d %H:%M:%S"))

    os.chdir(DOCKER_PATH)
    docker_client = docker.from_env()

    for service in stacks:
        service_dir = DOCKER_PATH / service
        if not service_dir.exists():
            logging.warning(f"Skipping {service} (directory does not exist)")
            continue

        logging.info(f"Updating service: {service}")
        try:
            run_compose_command(service_dir, ["pull"])
            run_compose_command(service_dir, ["up", "-d", "--build"])
            check_container_status(docker_client, service)
        except Exception as e:
            logging.error(f"Failed to update {service}: {e}")

    # Prune Docker system resources
    try:
        logging.info("Pruning unused containers, images, networks, and volumes...")
        docker_client.containers.prune()
        docker_client.images.prune()
        docker_client.networks.prune()
        docker_client.volumes.prune()
    except Exception as e:
        logging.error(f"Failed to prune system: {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Update Docker containers using Compose v2 CLI")
    parser.add_argument("-q", "--quiet", dest="stacks", nargs='*', default=DEFAULT_STACKS, help="List of compose services to recreate")

    args = parser.parse_args()
    stacks = args.stacks
    main(stacks)