# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.

import click
from rich.console import Console
from rich.table import Table

from llama_verifications.benchmarks.benchmarks.registry import BenchmarkRegistry


@click.command(name="list-benchmarks")
@click.option("--all", is_flag=True, help="Show all benchmarks including unverified ones")
def list_benchmarks_command(all):
    """List all available benchmarks."""
    console = Console()

    if all:
        table = Table(title="Available Benchmarks", title_style="bold yellow", show_lines=True)
        table.add_column("Benchmark ID", style="cyan")
        table.add_column("Development Status", style="white")
        table.add_column("Description", no_wrap=False, max_width=80)

        order = ["verified", "unverified", "under_development"]
        sorted_benchmarks = sorted(
            BenchmarkRegistry.get_benchmark_ids(),
            key=lambda x: (
                order.index(BenchmarkRegistry.get_benchmark(x).development_status),
                x,  # then alphabetically
            ),
        )
    else:
        table = Table(title="Verified Benchmarks", title_style="bold yellow", show_lines=True)
        table.add_column("Benchmark ID", style="cyan")
        table.add_column("Description", no_wrap=False, max_width=80)
        sorted_benchmarks = sorted(BenchmarkRegistry.get_verified_benchmark_ids())

    console.print("")

    for benchmark in sorted_benchmarks:
        bench = BenchmarkRegistry.get_benchmark(benchmark)
        if all:
            color = "magenta" if bench.development_status == "verified" else "white"
            table.add_row(
                benchmark,
                f"[{color}]{bench.development_status}[/]",
                bench.description,
            )
        else:
            table.add_row(benchmark, bench.description)

    console.print(table)

    if all:
        # show a note section with three bullets, each explaining what (verified, under_development, unverified) means
        console.print(
            """
[bold magenta]* verified[/]: The benchmark numbers have been cross-checked with internal or reported results.
[bold yellow]* unverified[/]: The benchmark is code-complete but dataset and numbers are not yet verified. Use with caution.
[bold red]* under_development[/]: The benchmark is still under development and may not be fully reliable. Do not use it.
            """
        )
