#!/usr/bin/env python3
"""
Docker Manager Application

This script manages Docker containers defined in subdirectories of `~/docker/`.
Each subdirectory is treated as a project and must contain a `docker-compose.yml`, `docker-compose.yaml`, `compose.yml`, or `compose.yaml`.

Usage Examples:
    1. List all running containers:
       $ dm list

    2. List only stacks with updates available:
       $ dm list --update

    3. Update a specific project (pulls latest images and recreates containers):
       $ dm update media-server

    3. Stop all projects (with confirmation prompt):
       $ dm stop

    4. Restart all projects without confirmation:
       $ dm restart --all

    5. Stop a specific project:
       $ dm stop web-app

Purpose:
    - Simplifies management of multiple Docker Compose projects.
    - Provides a unified interface for common operations (list, stop, update, restart).
    - Eliminates the need to navigate to specific directories to run docker commands.

To make the command available as dm from any directory, run:
  chmod +x ~/shell/docker-manager/docker-manager.py
  mkdir -p ~/.local/bin
  ln -s ~/shell/docker-manager/docker-manager.py ~/.local/bin/dm
"""

import os
import argparse
import subprocess
import json
from pathlib import Path

try:
    from rich.console import Console
    from rich.table import Table
    from rich.panel import Panel
    from rich import box
except ImportError:
    print("Error: The 'rich' library is required but not installed.")
    print("Please install it using: pip install -r ~/shell/docker-manager/requirements.txt")
    print("Or directly: pip install rich")
    exit(1)

# Configuration
DOCKER_ROOT = Path(os.path.expanduser("~/docker"))
console = Console()


def get_diun_container_name():
    """Find the running diun container name."""
    # Try to find by image name first (most reliable if standard image is used)
    cmd = ["docker", "ps", "--filter", "ancestor=crazymax/diun", "--format", "{{.Names}}"]
    res = subprocess.run(cmd, capture_output=True, text=True)
    if res.returncode == 0 and res.stdout.strip():
        return res.stdout.strip().split('\n')[0]

    # Fallback to finding by name containing 'diun'
    cmd = ["docker", "ps", "--filter", "name=diun", "--format", "{{.Names}}"]
    res = subprocess.run(cmd, capture_output=True, text=True)
    if res.returncode == 0 and res.stdout.strip():
        return res.stdout.strip().split('\n')[0]

    return None


def get_diun_info():
    """Fetch image information from Diun."""
    diun_info = {}
    try:
        diun_container = get_diun_container_name()

        if not diun_container:
            return diun_info

        # Get diun image list
        cmd = ["docker", "exec", diun_container, "diun", "image", "list", "--raw"]
        res = subprocess.run(cmd, capture_output=True, text=True)

        if res.returncode != 0:
            return diun_info

        data = json.loads(res.stdout)
        for img in data.get("images", []):
            name = img.get("name", "")
            latest = img.get("latest", {})
            digest = latest.get("digest", "")
            labels = latest.get("labels", {})

            # Normalize name for matching
            # Remove docker.io/ prefix
            norm_name = name
            if norm_name.startswith("docker.io/"):
                norm_name = norm_name[10:]

            # If no registry (no /), prepend library/
            if "/" not in norm_name:
                norm_name = f"library/{norm_name}"

            tag = latest.get("tag", "latest")

            if norm_name not in diun_info:
                diun_info[norm_name] = {}

            diun_info[norm_name][tag] = {
                "digest": digest,
                "labels": labels,
                "original_name": name
            }

    except (subprocess.CalledProcessError, json.JSONDecodeError):
        pass

    return diun_info


def get_image_version(labels):
    """Extract version from image labels."""
    version_keys = [
        "org.opencontainers.image.version",
        "version",
        "org.label-schema.version",
        "build_version"
    ]
    for key in version_keys:
        if key in labels:
            return labels[key]
    return "latest"


def get_projects():
    """Scan DOCKER_ROOT for subdirectories with compose files."""
    projects = {}
    if not DOCKER_ROOT.exists():
        print(f"Error: Directory {DOCKER_ROOT} does not exist.")
        return projects

    for item in DOCKER_ROOT.iterdir():
        if item.is_dir():
            if (item / "docker-compose.yml").exists() or \
               (item / "docker-compose.yaml").exists() or \
               (item / "compose.yml").exists() or \
               (item / "compose.yaml").exists():
                projects[item.name] = item
    return projects


def run_command(cmd, cwd, capture_output=False):
    """Run a shell command in a specific directory."""
    try:
        result = subprocess.run(
            cmd,
            cwd=cwd,
            check=True,
            text=True,
            capture_output=capture_output
        )
        return result
    except subprocess.CalledProcessError as e:
        if capture_output:
            return e
        print(f"Error running command: {' '.join(cmd)}")
        print(f"cwd: {cwd}")
        return None


def list_containers(projects, show_updates_only=False):
    """List running containers for all projects."""
    table = Table(title="Docker Containers", box=box.ROUNDED)
    table.add_column("Project", style="cyan", no_wrap=True)
    table.add_column("Container Name", style="green")
    table.add_column("State", style="yellow")
    table.add_column("Image", style="blue")
    table.add_column("Version", style="magenta")
    table.add_column("Update", style="red")

    # Get Diun info once
    diun_info = get_diun_info()

    found_any = False
    for name, path in sorted(projects.items()):

        # Buffer for project rows
        project_rows = []
        project_has_update = False

        # Get running container names
        cmd = ["docker", "compose", "ps", "--format", "json"]
        res = run_command(cmd, path, capture_output=True)

        if res and res.returncode == 0:
            lines = res.stdout.strip().split('\n')
            for line in lines:
                if not line:
                    continue
                try:
                    container = json.loads(line)
                    c_name = container.get('Name', '')
                    state = container.get('State', '')
                    image = container.get('Image', '')

                    if not c_name:
                        continue

                    # Get version and update status
                    version = "unknown"
                    update_status = ""

                    # Normalize image name for lookup
                    norm_image = image
                    tag = "latest"

                    # Remove tag if present
                    if ":" in norm_image:
                        base, sep, t = norm_image.rpartition(":")
                        if sep and "/" not in t:
                            norm_image = base
                            tag = t

                    if norm_image.startswith("docker.io/"):
                        norm_image = norm_image[10:]
                    if "/" not in norm_image:
                        norm_image = f"library/{norm_image}"

                    # Check Diun info
                    if diun_info:
                        if norm_image in diun_info:
                            image_tags = diun_info[norm_image]

                            if tag in image_tags:
                                info = image_tags[tag]
                                version = get_image_version(info.get("labels", {}))

                                inspect_id_cmd = ["docker", "inspect", c_name, "--format", "{{.Image}}"]
                                inspect_id_res = run_command(inspect_id_cmd, path, capture_output=True)

                                if inspect_id_res and inspect_id_res.returncode == 0:
                                    image_id = inspect_id_res.stdout.strip()

                                    inspect_digest_cmd = ["docker", "inspect", image_id, "--format", "{{if .RepoDigests}}{{index .RepoDigests 0}}{{end}}"]
                                    inspect_digest_res = run_command(inspect_digest_cmd, path, capture_output=True)

                                    if inspect_digest_res and inspect_digest_res.returncode == 0:
                                        running_digest_full = inspect_digest_res.stdout.strip()
                                        if "@" in running_digest_full:
                                            running_digest = running_digest_full.split("@")[1]
                                            latest_digest = info.get("digest", "")

                                            if latest_digest:
                                                if running_digest != latest_digest:
                                                    update_status = "[bold red]Update Available[/bold red]"
                                                    project_has_update = True
                                                else:
                                                    update_status = "[green]Up to Date[/green]"
                            else:
                                update_status = "[dim]Tag Not Monitored[/dim]"
                        else:
                            update_status = "[dim]Not Monitored[/dim]"

                    if version in ("latest", "unknown"):
                        inspect_cmd = ["docker", "inspect", c_name, "--format", "{{json .Config.Labels}}"]
                        inspect_res = run_command(inspect_cmd, path, capture_output=True)
                        if inspect_res and inspect_res.returncode == 0:
                            try:
                                labels = json.loads(inspect_res.stdout)
                                version = get_image_version(labels)
                            except Exception:
                                pass

                    project_rows.append((name, c_name, state, image, version, update_status))

                except json.JSONDecodeError:
                    pass

        # If not hiding, OR if update, add to table
        if not show_updates_only or project_has_update:
            if project_rows:
                if found_any:
                    table.add_section()
                found_any = True
                for row in project_rows:
                    table.add_row(*row)

    if found_any:
        console.print(table)
    else:
        if show_updates_only:
            console.print("[green]No updates available for running containers.[/green]")
        else:
            console.print("[yellow]No running containers found in managed projects.[/yellow]")


def is_project_running(path):
    """Check if any containers in the project are running."""
    cmd = ["docker", "compose", "ps", "--format",
           "json", "--filter", "status=running"]
    res = run_command(cmd, path, capture_output=True)
    if res and res.returncode == 0:
        # If output is not empty/just brackets, something is running
        output = res.stdout.strip()
        return output and output != "[]"
    return False


def describe_project(projects, target):
    """Show detailed information about containers in a project."""
    if target not in projects:
        console.print(f"[red]Error: Project '{target}' not found.[/red]")
        return

    path = projects[target]
    console.print(Panel(f"[bold blue]Describing project: {target}[/bold blue]", expand=False))

    cmd = ["docker", "compose", "ps", "--format", "json"]
    res = run_command(cmd, path, capture_output=True)

    if res and res.returncode == 0:
        lines = res.stdout.strip().split('\n')
        found_any = False
        for line in lines:
            if not line:
                continue
            try:
                container = json.loads(line)
                found_any = True
                name = container.get('Name', 'N/A')

                # Create a table for this container
                table = Table(show_header=False, box=box.SIMPLE)
                table.add_column("Field", style="bold cyan")
                table.add_column("Value")

                table.add_row("Service", container.get('Service', 'N/A'))
                table.add_row("Container", name)
                table.add_row("Image", container.get('Image', 'N/A'))
                table.add_row("State", container.get('State', 'N/A'))
                table.add_row("Status", container.get('Status', 'N/A'))

                ports = container.get('Ports', '')
                if ports:
                    table.add_row("Ports", ports)

                # Get details via inspect
                inspect_cmd = ["docker", "inspect", name, "--format", "{{json .}}"]
                inspect_res = run_command(inspect_cmd, path, capture_output=True)
                if inspect_res and inspect_res.returncode == 0:
                    try:
                        details = json.loads(inspect_res.stdout.strip())

                        # Description
                        labels = details.get("Config", {}).get("Labels", {})
                        description = labels.get("org.opencontainers.image.description")
                        if description:
                            table.add_row("Description", description)

                        console.print(Panel(table, title=f"[bold]{name}[/bold]", border_style="green"))

                        # Volumes/Mounts
                        mounts = details.get("Mounts", [])
                        if mounts:
                            v_table = Table(title="Volumes", box=box.MINIMAL)
                            v_table.add_column("Type", style="cyan")
                            v_table.add_column("Source", style="yellow")
                            v_table.add_column("Destination", style="green")

                            for mount in mounts:
                                m_type = mount.get("Type")
                                source = mount.get("Source", "")
                                dest = mount.get("Destination", "")
                                v_table.add_row(m_type, source, dest)

                            console.print(v_table)
                            console.print("") # Spacing

                    except json.JSONDecodeError:
                        pass
                else:
                    console.print(Panel(table, title=f"[bold]{name}[/bold]", border_style="green"))

            except json.JSONDecodeError:
                pass

        if not found_any:
            console.print("[yellow]No containers found for this project.[/yellow]")
    else:
        console.print("[red]Failed to get container info.[/red]")


def list_project_volumes(projects, target):
    """List volumes used by containers in a project."""
    if target not in projects:
        console.print(f"[red]Error: Project '{target}' not found.[/red]")
        return

    path = projects[target]

    table = Table(title=f"Volumes for project: {target}", box=box.ROUNDED)
    table.add_column("Service", style="cyan")
    table.add_column("Container", style="blue")
    table.add_column("Type", style="magenta")
    table.add_column("Source", style="yellow")
    table.add_column("Destination", style="green")

    cmd = ["docker", "compose", "ps", "--format", "json"]
    res = run_command(cmd, path, capture_output=True)

    if res and res.returncode == 0:
        lines = res.stdout.strip().split('\n')
        for line in lines:
            if not line:
                continue
            try:
                container = json.loads(line)
                name = container.get('Name', 'N/A')
                service = container.get('Service', 'N/A')

                inspect_cmd = ["docker", "inspect", name, "--format", "{{json .Mounts}}"]
                inspect_res = run_command(inspect_cmd, path, capture_output=True)

                if inspect_res and inspect_res.returncode == 0:
                    try:
                        mounts = json.loads(inspect_res.stdout.strip())
                        if mounts:
                            for mount in mounts:
                                m_type = mount.get("Type", "N/A")
                                source = mount.get("Source", "N/A")
                                dest = mount.get("Destination", "N/A")
                                table.add_row(service, name, m_type, source, dest)
                    except json.JSONDecodeError:
                        pass
            except json.JSONDecodeError:
                pass

    console.print(table)



def manage_project(projects, action, target, extra_args=None, container_filter=None):
    """
    Execute the specified action (stop, update, restart, logs) on target project(s).

    Args:
        projects (dict): Dictionary of project names to their paths.
        action (str): The action to perform ('stop', 'update', 'restart', 'logs').
        target (str): The target project name or 'all'.
        extra_args (list): Additional arguments for the command (e.g. for logs).
        container_filter (str): Optional container/service name to filter logs.
    """

    targets = []
    # Determine which projects to target
    if target == "all":
        targets = sorted(projects.items())
    elif target in projects:
        targets = [(target, projects[target])]
    else:
        console.print(f"[red]Error: Project '{target}' not found in {DOCKER_ROOT}[/red]")
        console.print(f"Available projects: {', '.join(sorted(projects.keys()))}")
        return

    # Logs are special, usually run on one project interactively
    if action == "logs":
        if len(targets) > 1:
            console.print("[red]Error: Logs can only be viewed for one project at a time.[/red]")
            return
        name, path = targets[0]

        # Filter by container/service if specified
        services_to_log = []
        if container_filter:
            cmd_ps = ["docker", "compose", "ps", "--format", "json"]
            res = run_command(cmd_ps, path, capture_output=True)
            if res and res.returncode == 0:
                try:
                    lines = res.stdout.strip().split('\n')
                    for line in lines:
                        if not line: continue
                        c = json.loads(line)
                        s_name = c.get('Service', '')
                        c_name = c.get('Name', '')
                        # Check if filter matches service or container name
                        if container_filter in s_name or container_filter in c_name:
                            services_to_log.append(s_name)
                except json.JSONDecodeError:
                    pass

            if not services_to_log:
                console.print(f"[yellow]No matching service/container found for '{container_filter}'. Showing all logs.[/yellow]")
            else:
                services_to_log = list(set(services_to_log))
                console.print(f"[blue]Showing logs for service(s): {', '.join(services_to_log)}[/blue]")

        console.print(f"[blue]Viewing logs for {name}...[/blue]")
        cmd = ["docker", "compose", "logs"] + (extra_args or [])

        if services_to_log:
            cmd.extend(services_to_log)

        # For logs, we want to stream output directly to stdout, so no capture_output
        # and we might want to allow Ctrl+C to exit gracefully
        try:
            subprocess.run(cmd, cwd=path, check=False)
        except KeyboardInterrupt:
            console.print("\n[yellow]Log view stopped.[/yellow]")
        return

    console.print(f"[bold]Performing '{action}' on {len(targets)} project(s)...[/bold]")

    for name, path in targets:
        console.print(f"\n[bold cyan][{name}] -> {action}...[/bold cyan]")

        # Execute the requested action
        if action == "stop":
            # Stop containers without removing them
            run_command(["docker", "compose", "stop"], path)
        elif action == "restart":
            # Restart containers
            run_command(["docker", "compose", "restart"], path)
        elif action == "update":
            # Update process: Check running -> Stop -> Pull -> Start if was running
            was_running = is_project_running(path)

            if was_running:
                console.print(f"  [yellow]Stopping running containers for {name}...[/yellow]")
                run_command(["docker", "compose", "stop"], path)

            console.print(f"  [blue]Pulling latest images for {name}...[/blue]")
            pull_res = run_command(["docker", "compose", "pull"], path)

            if pull_res and pull_res.returncode == 0:
                if was_running:
                    console.print(f"  [green]Restarting containers for {name}...[/green]")
                    run_command(["docker", "compose", "up", "-d"], path)
                else:
                    console.print(
                        f"  [green]Project {name} was not running. Images updated, but not started.[/green]")
            else:
                console.print(f"  [red]Failed to pull images for {name}. Skipping update.[/red]")


def main():
    """Main entry point for the application."""
    parser = argparse.ArgumentParser(
        description="Manage Docker containers in ~/docker/")
    subparsers = parser.add_subparsers(
        dest="command", help="Command to execute")

    # List command configuration
    list_parser = subparsers.add_parser("list", help="List running containers")
    list_parser.add_argument("--update", action="store_true", help="Show only stacks that have updates available")

    # Describe command configuration
    describe_parser = subparsers.add_parser("describe", help="Show details of a project's containers")
    describe_parser.add_argument("project", help="Project name to describe")

    # Volumes command configuration
    volumes_parser = subparsers.add_parser("volumes", help="List volumes used by a project")
    volumes_parser.add_argument("project", help="Project name to list volumes for")

    # Stop command configuration
    stop_parser = subparsers.add_parser("stop", help="Stop containers")
    stop_parser.add_argument("project", nargs="?", default=None,
                             help="Project name (optional). If omitted, asks for confirmation to stop ALL.")
    stop_parser.add_argument("--all", action="store_true",
                             help="Stop all projects without confirmation prompt if specified")

    # Update command configuration
    update_parser = subparsers.add_parser(
        "update", help="Pull and update containers")
    update_parser.add_argument("project", nargs="?", default=None,
                               help="Project name (optional). If omitted, asks for confirmation to update ALL.")
    update_parser.add_argument("--all", action="store_true",
                               help="Update all projects without confirmation prompt if specified")

    # Restart command configuration
    restart_parser = subparsers.add_parser(
        "restart", help="Restart containers")
    restart_parser.add_argument("project", nargs="?", default=None,
                                help="Project name (optional). If omitted, asks for confirmation to restart ALL.")
    restart_parser.add_argument("--all", action="store_true",
                                help="Restart all projects without confirmation prompt if specified")

    # Logs command configuration
    logs_parser = subparsers.add_parser("logs", help="View container logs")
    logs_parser.add_argument("project", help="Project name to view logs for")
    logs_parser.add_argument("container", nargs="?", help="Container/Service name to filter logs (optional)")
    logs_parser.add_argument(
        "-f", "--follow", action="store_true", help="Follow log output")
    logs_parser.add_argument(
        "--tail", default="all", help="Number of lines to show from the end of the logs (default: all)")

    args = parser.parse_args()

    # Scan for projects
    projects = get_projects()
    if not projects:
        print(f"No docker projects found in {DOCKER_ROOT}")
        return

    # Dispatch commands
    if args.command == "list":
        list_containers(projects, show_updates_only=args.update)
    elif args.command == "describe":
        describe_project(projects, args.project)
    elif args.command == "volumes":
        list_project_volumes(projects, args.project)
    elif args.command == "logs":
        # Prepare extra args for logs
        extra_args = []
        if args.follow:
            extra_args.append("-f")
        if args.tail:
            extra_args.extend(["--tail", args.tail])

        manage_project(projects, "logs", args.project, extra_args=extra_args, container_filter=args.container)

    elif args.command in ["stop", "update", "restart"]:
        target = args.project

        # Handle "all" logic safely with user confirmation
        if target is None:
            if args.all:
                target = "all"
            else:
                # Interactive check if user wants to apply to all
                print(
                    f"No project specified. Do you want to {args.command} ALL projects? (y/N)")
                choice = input().lower()
                if choice == 'y':
                    target = "all"
                else:
                    print("Operation cancelled. Specify a project name or use --all.")
                    return

        manage_project(projects, args.command, target)
    else:
        # Show help if no command provided
        parser.print_help()


if __name__ == "__main__":
    try:
        main()
    finally:
        console.show_cursor(True)
