from math import ceil

from rich import print
from rich.table import Table
from rich.panel import Panel
from rich.console import Group, Console
from rich.align import Align


from rlmesh.client.console import RLMeshConsole
from rlmesh.proto import cp_public_pb2


def _into_columns(items, cols):
    if cols < 1:
        cols = 1
    rows = ceil(len(items) / cols)
    out = []
    for c in range(cols):
        start = c * rows
        out.append(items[start : start + rows])
    return out, rows


def print_envs(res: cp_public_pb2.ListEnvironmentsRes):
    console = Console()
    panels = []
    for pkg, envs in res.packages.items():
        env_list = list(envs.env_ids)
        if not env_list:
            continue

        max_len = max(len(s) for s in env_list)
        term_width = console.size.width or 100
        min_col_width = max(18, min(40, max_len + 2))
        cols = max(1, min(4, term_width // min_col_width))

        cols_data, rows = _into_columns(env_list, cols)

        grid = Table.grid(expand=True, pad_edge=False)
        for _ in range(cols):
            grid.add_column(ratio=1, justify="left", no_wrap=True)

        for r in range(rows):
            grid.add_row(
                *[
                    (cols_data[c][r] if r < len(cols_data[c]) else "")
                    for c in range(cols)
                ]
            )

        panels.append(
            Panel(
                grid,
                title=f"[green]{pkg}[/green]",
                border_style="bright_black",
                padding=(0, 1),
            )
        )

    group = Group(*panels)
    print(
        Panel(
            group,
            title="[bold yellow]Available Environments[/bold yellow]",
            border_style="bold yellow",
            expand=True,
        )
    )


def generate_live_env_panel(
    console: RLMeshConsole, step_count, num_envs, active_envs, score, fps, elapsed, perf
):
    # env_ms = perf.get("env_step_ms", 0)
    rtt_ms = perf.get("downstream_rtt_ms", 0)

    total_ms = rtt_ms
    theoretical_fps = 1000.0 / total_ms if total_ms > 0 else 0

    utilization = (fps / theoretical_fps * 100) if theoretical_fps > 0 else 0
    # if utilization > 70:
    #     status = "[green]Optimal[/green]"
    # elif utilization > 50:
    #     status = "[yellow]Network limited[/yellow]"
    # elif theoretical_fps == 0:
    #     status = "[dim]Measuring...[/dim]"
    # else:
    #     status = "[red]Severely limited[/red]"

    header = console.table().grid(expand=True, padding=2)
    header.add_row(
        f"[cyan]Steps:[/cyan] {step_count}",
        f"[green]Active:[/green] {active_envs}/{num_envs}",
        f"[magenta]Score:[/magenta] {score:.2f}",
        f"[blue]FPS:[/blue] {fps:.1f}",
        f"[yellow]Time:[/yellow] {elapsed:.1f}s",
    )

    perf_tbl = console.table().grid(expand=False, padding=(0, 2))
    perf_tbl.add_row("Theoretical Max", f"{theoretical_fps:.1f} steps/s")
    perf_tbl.add_row("Actual Observed", f"{fps:.1f} steps/s")
    perf_tbl.add_row("Utilization", f"{utilization:.1f}%")
    # perf_tbl.add_row("Status", status)

    grid = console.table().grid(expand=True)
    grid.add_row(Align.center(header))
    grid.add_row("")
    grid.add_row(Align.center(perf_tbl))

    return console.panel(
        Align.center(grid),
        title="[bold white]Environment Performance[/bold white]",
    )


def final_results_panel(console, results: dict):
    table = console.table(box=None, show_header=False, expand=False)
    table.add_row("[cyan]Final Score[/cyan]", f"{results['score']:.4f}")
    table.add_row("[green]Duration (s)[/green]", f"{results['duration']:.2f}")
    table.add_row("[yellow]Timestamp[/yellow]", results["timestamp"])

    return console.panel(
        Align.left(table),
        title="[bold white]Session Results[/bold white]",
    )
