freepik/freepik_cli/utils/console.py

"""Rich console singleton, theme, banner, and all display helpers."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import rich.box
from rich.align import Align
from rich.columns import Columns
from rich.console import Console
from rich.panel import Panel
from rich.rule import Rule
from rich.syntax import Syntax
from rich.table import Table
from rich.text import Text
from rich.theme import Theme

FREEPIK_THEME = Theme(
    {
        "info": "bold cyan",
        "success": "bold green",
        "warning": "bold yellow",
        "error": "bold red",
        "model": "bold magenta",
        "taskid": "bold blue",
        "path": "bold white underline",
        "dim.label": "dim white",
        "highlight": "bold magenta",
        "brand": "bold magenta",
    }
)

console = Console(theme=FREEPIK_THEME, highlight=True)
err_console = Console(stderr=True, theme=FREEPIK_THEME)

BANNER = """\
[bold magenta] ███████╗██████╗ ███████╗███████╗██████╗ ██╗██╗  ██╗[/bold magenta]
[bold magenta] ██╔════╝██╔══██╗██╔════╝██╔════╝██╔══██╗██║██║ ██╔╝[/bold magenta]
[bold magenta] █████╗  ██████╔╝█████╗  █████╗  ██████╔╝██║█████╔╝ [/bold magenta]
[bold magenta] ██╔══╝  ██╔══██╗██╔══╝  ██╔══╝  ██╔═══╝ ██║██╔═██╗ [/bold magenta]
[bold magenta] ██║     ██║  ██║███████╗███████╗██║     ██║██║  ██╗[/bold magenta]
[bold magenta] ╚═╝     ╚═╝  ╚═╝╚══════╝╚══════╝╚═╝     ╚═╝╚═╝  ╚═╝[/bold magenta]
[dim]            AI Media Generation CLI  •  v0.1.0[/dim]"""


def print_banner() -> None:
    console.print()
    console.print(Align.center(Text.from_markup(BANNER)))
    console.print(Rule(style="magenta dim"))
    console.print()


@dataclass
class GenerationResult:
    task_id: str
    model: str
    output_path: Path
    width: Optional[int] = None
    height: Optional[int] = None
    seed: Optional[int] = None
    duration: Optional[str] = None
    task_type: str = "image"


def print_result(result: GenerationResult) -> None:
    table = Table(
        show_header=False,
        box=rich.box.SIMPLE,
        padding=(0, 1),
        show_edge=False,
    )
    table.add_column(style="dim.label", width=16, no_wrap=True)
    table.add_column(style="bold", overflow="fold")

    table.add_row("Model", f"[model]{result.model}[/model]")
    table.add_row("Task ID", f"[taskid]{result.task_id[:16]}…[/taskid]")

    if result.width and result.height:
        table.add_row("Dimensions", f"{result.width} × {result.height} px")
    if result.duration:
        table.add_row("Duration", f"{result.duration}s")
    if result.seed is not None:
        table.add_row("Seed", str(result.seed))
    else:
        table.add_row("Seed", "[dim]random[/dim]")

    table.add_row("Saved to", f"[path]{result.output_path}[/path]")

    title_map = {
        "image": "[success] Image Generated [/success]",
        "video": "[success] Video Generated [/success]",
        "upscale-image": "[success] Image Upscaled [/success]",
        "upscale-video": "[success] Video Upscaled [/success]",
        "icon": "[success] Icon Generated [/success]",
        "expand": "[success] Image Expanded [/success]",
        "describe": "[success] Image Described [/success]",
    }
    title = title_map.get(result.task_type, "[success] Done [/success]")

    console.print(Panel(table, title=title, border_style="green", padding=(1, 2)))


def print_describe_result(task_id: str, prompt_text: str, output_path: Optional[Path] = None) -> None:
    content = Text(prompt_text, style="italic")
    footer = ""
    if output_path:
        footer = f"\n\n[dim]Saved to:[/dim] [path]{output_path}[/path]"

    console.print(
        Panel(
            Text.from_markup(f"{prompt_text}{footer}"),
            title="[success] Image Description [/success]",
            subtitle=f"[dim]Task: {task_id[:16]}…[/dim]",
            border_style="green",
            padding=(1, 2),
        )
    )


def print_no_wait(task_id: str, task_type: str, model: str) -> None:
    console.print(
        Panel(
            Text.from_markup(
                f"[info]Task submitted successfully.[/info]\n\n"
                f"[dim.label]Task ID:[/dim.label]  [taskid]{task_id}[/taskid]\n"
                f"[dim.label]Type:[/dim.label]     {task_type}\n"
                f"[dim.label]Model:[/dim.label]    [model]{model}[/model]\n\n"
                f"[dim]The task is processing asynchronously. Results will be available\n"
                f"on the Freepik dashboard when complete.[/dim]"
            ),
            title="⏳ Task Queued",
            border_style="cyan",
            padding=(1, 2),
        )
    )


def print_error(message: str, hint: Optional[str] = None) -> None:
    body = f"[error]{message}[/error]"
    if hint:
        body += f"\n\n[dim]Hint:[/dim] {hint}"
    err_console.print(
        Panel(
            Text.from_markup(body),
            title="[error] Error [/error]",
            border_style="red",
            padding=(1, 2),
        )
    )


def print_warning(message: str) -> None:
    console.print(
        Panel(
            Text.from_markup(f"[warning]{message}[/warning]"),
            title="[warning] Warning [/warning]",
            border_style="yellow",
            padding=(0, 2),
        )
    )


def print_config_table(config_dict: dict, masked_keys: set[str] | None = None) -> None:
    masked_keys = masked_keys or {"api_key"}
    table = Table(
        title="[brand]Freepik CLI Configuration[/brand]",
        show_header=True,
        header_style="bold magenta",
        border_style="magenta",
        box=rich.box.ROUNDED,
        padding=(0, 1),
    )
    table.add_column("Key", style="dim.label", width=28)
    table.add_column("Value", style="bold", overflow="fold")

    for key, value in config_dict.items():
        if key in masked_keys and value:
            display = f"[dim]{'*' * 8}{str(value)[-4:]}[/dim]" if len(str(value)) > 4 else "[dim]****[/dim]"
        elif value is None or value == "":
            display = "[dim]not set[/dim]"
        else:
            display = str(value)
        table.add_row(key, display)

    console.print(table)


def print_config_toml(config_dict: dict, masked_keys: set[str] | None = None) -> None:
    masked_keys = masked_keys or {"api_key"}
    lines = []
    for key, value in config_dict.items():
        if key in masked_keys:
            continue
        if value is None:
            continue
        if isinstance(value, str):
            lines.append(f'{key} = "{value}"')
        elif isinstance(value, bool):
            lines.append(f"{key} = {str(value).lower()}")
        else:
            lines.append(f"{key} = {value}")

    toml_str = "\n".join(lines)
    console.print(
        Panel(
            Syntax(toml_str, "toml", theme="monokai", background_color="default"),
            title="[brand]~/.config/freepik-cli/config.toml[/brand]",
            border_style="magenta",
            padding=(1, 2),
        )
    )