"""AI Infrastructure Service CLI - Cluster info."""

from enum import Enum
from typing import Annotated

import typer
from msgspec import to_builtins
from msgspec.json import encode as json_encode
from msgspec.yaml import encode as yaml_encode
from rich.table import Table
from slices_clientlib_aiis.schema.worker import GpuModel

from . import console
from .ai import ai_app
from .ai_context import AIContext
from .ai_helpers import (
    handle_some_exceptions,
)


class ListClustersFormat(Enum):
    """List Clusters Formats."""

    table = "table"
    json = "json"
    yaml = "yaml"


@ai_app.command("clusters", rich_help_panel="Information commands")
@handle_some_exceptions
def list_clusters(
    ctx: AIContext,
    cluster_id: Annotated[str | None, typer.Argument(help="ID of the cluster to show info for.")] = None,
    all: Annotated[bool, typer.Option("--all", "-a", help="(if no <cluster_id>) Show all machine info.")] = False,
    gpu_details: Annotated[bool, typer.Option("--gpu-details", help="Show more details about GPUs")] = False,
    format: Annotated[
        ListClustersFormat,
        typer.Option("--format", "-f", help="Format in which the resources should be listed."),
    ] = "table",  # type: ignore
):
    """Show available clusters and their resources."""

    cluster_info = ctx.obj.api.for_project(ctx.obj.require_project_id()).ai(ctx.obj.site_id).workers.list_clusters()

    if format in (ListClustersFormat.json, ListClustersFormat.yaml):
        if cluster_id or all:
            # List one or all clusters, with machine info
            machine_info = (
                ctx.obj.api.for_project(ctx.obj.require_project_id()).ai(ctx.obj.site_id).workers.list_slaves()
            )

            clusters_with_machines = []
            for cluster in cluster_info:
                if cluster_id and str(cluster.cluster_id) != str(cluster_id):
                    continue

                cluster_dict = to_builtins(cluster)
                cluster_dict["machines"] = [
                    to_builtins(machine)
                    for machine in machine_info
                    if str(machine.cluster_id) == str(cluster.cluster_id)
                ]
                clusters_with_machines.append(cluster_dict)

            if format == ListClustersFormat.json:
                console.print_json(json_encode(clusters_with_machines).decode())
            elif format == ListClustersFormat.yaml:
                console.print(yaml_encode(clusters_with_machines).decode())
        elif format == ListClustersFormat.json:
            console.print_json(json_encode(list(cluster_info)).decode())
        elif format == ListClustersFormat.yaml:
            console.print(yaml_encode(list(cluster_info)).decode())
        return

    def _gpu_str(gpu: GpuModel) -> str:
        if gpu_details:
            return f"{gpu.name} {gpu.vendor} ({gpu.memory_mb / 1024.0:.2f}GB)"
        return gpu.name

    if cluster_id or all:
        # List one or all clusters, with machine info
        machine_info = ctx.obj.api.for_project(ctx.obj.require_project_id()).ai(ctx.obj.site_id).workers.list_slaves()

        for cluster in cluster_info:
            if cluster_id and str(cluster.cluster_id) != str(cluster_id):
                continue

            console.print(f"[bold]Cluster {cluster.cluster_id}[/bold] - {cluster.comment}")
            cols = [
                "Machine",
                "GPU Model",
                "GPUs",
                "CPUs",
                "Memory (GB)",
                "Accepting Jobs",
                "Active Jobs",
                "CUDA version",
                "CPU model",
            ]
            table = Table(*cols, title=f"Machines in Cluster {cluster.cluster_id}")
            for machine in machine_info:
                if str(machine.cluster_id) == str(cluster.cluster_id):
                    table.add_row(
                        machine.name,
                        ",".join([_gpu_str(g) for g in cluster.gpu_model]),
                        f"{machine.gpu.available}/{machine.gpu.acquired}",
                        f"{machine.cpu.available}/{machine.cpu.acquired}",
                        f"{machine.cpu_memory_mb.available / 1000:.2f}/{machine.cpu_memory_mb.acquired / 1000:.2f}",
                        str(machine.accepting_jobs),
                        str(len(machine.active_job_uuids)) if machine.active_job_uuids is not None else "-",
                        machine.cuda_version_full,
                        ",".join(machine.cpu_model),
                    )
            console.print(table)
    else:
        # List all clusters, without any machine info
        cols = ["ID", "GPU Model", "Comment", "Machines (count)", "GPUs", "CPUs", "Memory (GB)"]

        table = Table(*cols, title="Clusters")
        for cluster in cluster_info:
            table.add_row(
                str(cluster.cluster_id),
                ",".join([_gpu_str(g) for g in cluster.gpu_model]),
                cluster.comment,
                str(cluster.slave_count),
                f"{cluster.gpu.available}/{cluster.gpu.acquired}",
                f"{cluster.cpu.available}/{cluster.cpu.acquired}",
                f"{cluster.cpu_memory_mb.available / 1000:.2f}/{cluster.cpu_memory_mb.acquired / 1000:.2f}",
            )
        console.print(table)
