#!/usr/bin/env python3

"""![status](/badges/status-cli-main.svg) Typer-based CLI for evaLLM evaluation.

Provides commands for running evaluations and inspecting available metrics.
"""

from __future__ import annotations

import sys
from collections.abc import Callable, Iterable
from pathlib import Path
from time import perf_counter
from typing import Annotated, Any, BinaryIO, Literal, NoReturn, cast

import orjson
import typer
import yaml
from pydantic import ValidationError
from rich.console import Console, Group
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
from rich.syntax import Syntax
from rich.table import Table
from rich.text import Text

from evallm.application.api import (
    count_characters,
    count_encodings,
    count_tokens,
)
from evallm.application.api import (
    metric as metric_spec,
)
from evallm.application.errors import (
    DatasetAccessError,
    DuplicateMetricResultKeyError,
    MetricConfigurationError,
    UnknownMetricError,
    UnsupportedDatasetError,
)
from evallm.application.models import EvaluationRequest, MetricSpec
from evallm.application.services import Evaluator
from evallm.metrics import builtin as builtin_metrics  # noqa: F401
from evallm.metrics.registry import MetricsCatalog

app = typer.Typer(
    no_args_is_help=True,
    suggest_commands=True,
    help="[bold]evaLLM[/bold]: Command-line interface.",
    epilog="For in-depth usage examples, see the [link=https://psandhaas.github.io/evaLLM/cli/]evaLLM CLI reference[/link].",
)
metrics_app = typer.Typer(help="Inspect available metrics and presets.")
app.add_typer(metrics_app, name="metrics")
err_console = Console(stderr=True)
out_console = Console()

_PRESET_FACTORIES: dict[str, Callable[[], MetricSpec]] = {
    "count_characters": count_characters,
    "count_tokens": count_tokens,
}


def _exit(message: str) -> NoReturn:
    """Print a styled error message and exit the CLI.

    Args:
        message (str): Error message to display.

    Raises:
        typer.Exit: Always raised with exit code 1.
    """
    err_console.rule(
        "[bold red]:exclamation: Error :exclamation:[/bold red]", style="red rule.line"
    )
    err_console.print(message, style="bold red", justify="center")
    raise typer.Exit(code=1)


def _load_mapping_file(path: Path) -> dict[str, Any]:
    """Load a JSON or YAML mapping from disk.

    Args:
        path (Path): File path to read.

    Returns:
        dict[str, Any]: Parsed mapping content.

    Raises:
        typer.Exit: If file type is unsupported or content is not a mapping.
    """  # noqa: DOC502
    suffix = path.suffix.lower()
    if suffix == ".json":
        data = orjson.loads(path.read_bytes())
    elif suffix in {".yaml", ".yml"}:
        data = yaml.safe_load(path.read_text(encoding="utf-8"))
    else:
        _exit(f"Unsupported file type for {path}. Use .json, .yaml, or .yml.")

    if not isinstance(data, dict):
        _exit(f"{path} must contain a JSON/YAML object at the top level.")
    return cast(dict[str, Any], data)


def _parse_metric_spec(value: str) -> MetricSpec:
    """Parse a metric argument into a MetricSpec.

    Args:
        value (str): Metric value from CLI input.

    Returns:
        MetricSpec: Parsed metric specification.

    Raises:
        typer.Exit: If the value cannot be parsed or validated.
    """  # noqa: DOC502
    if value in _PRESET_FACTORIES:
        return _PRESET_FACTORIES[value]()

    if value.startswith("count_tokens:"):
        segmentation = value.split(":", 1)[1].strip().lower()
        if segmentation not in {"whitespace", "wordpunct"}:
            _exit(
                "count_tokens preset expects segmentation in {whitespace, wordpunct}."
            )
        return count_tokens(
            segmentation=cast(Literal["whitespace", "wordpunct"], segmentation)
        )

    if value.startswith("count_encodings:"):
        model = value.split(":", 1)[1].strip()
        if not model:
            _exit(
                "count_encodings preset expects a model name, e.g. count_encodings:bert-base-uncased"
            )
        return count_encodings(model)

    candidate = Path(value)
    if candidate.exists():
        raw: Any = _load_mapping_file(candidate)
    else:
        try:
            raw = orjson.loads(value)
        except orjson.JSONDecodeError:
            return metric_spec(value)

    if isinstance(raw, str):
        return metric_spec(raw)
    if not isinstance(raw, dict):
        _exit(
            "--metric must be a metric name, preset, JSON object, or path to JSON/YAML."
        )

    try:
        return MetricSpec.model_validate(raw)
    except ValidationError as exc:
        _exit(f"Invalid metric spec '{value}': {exc}")


def _load_run_config(config_path: Path | None) -> dict[str, Any]:
    """Load run configuration from an optional file path.

    Args:
        config_path (Path | None): Optional path to a JSON or YAML config file.

    Returns:
        dict[str, Any]: Parsed configuration mapping, or an empty mapping if no path is
            provided.

    Raises:
        typer.Exit: If the config path does not exist or is invalid.
    """  # noqa: DOC502
    if config_path is None:
        return {}
    if not config_path.exists():
        _exit(f"Config file not found: {config_path}")
    return _load_mapping_file(config_path)


def _open_output(path: Path | None) -> tuple[BinaryIO, bool]:
    """Resolve the output stream and ownership flag.

    Args:
        path (Path | None): Optional output path. None or '-' maps to stdout.

    Returns:
        tuple[BinaryIO, bool]: A tuple of (binary stream, should_close flag).
    """
    if path is None or str(path) == "-":
        return sys.stdout.buffer, False
    return path.open("wb"), True


def _stream_results(evaluator: Evaluator, output: BinaryIO) -> tuple[int, float]:
    """Stream evaluation results to output with progress feedback.

    Args:
        evaluator (Evaluator): Evaluator instance to execute.
        output (BinaryIO): Binary output stream for JSONL results.

    Returns:
        tuple[int, float]: A tuple of (processed record count, elapsed seconds).
    """
    count = 0
    started = perf_counter()

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        TextColumn("•"),
        TimeElapsedColumn(),
        console=err_console,
        transient=True,
    ) as progress:
        task = progress.add_task("Processing... 0 items", total=None)

        for result in evaluator.run():
            output.write(orjson.dumps(result.model_dump()) + b"\n")
            output.flush()
            count += 1

            if count == 1 or count % 25 == 0:
                progress.update(task, description=f"Processing... {count} items")

    elapsed_s = perf_counter() - started
    return count, elapsed_s


def _as_metric_dicts(metrics: Iterable[MetricSpec]) -> list[dict[str, Any]]:
    """Convert metric specs into serializable dictionaries.

    Args:
        metrics (Iterable[MetricSpec]): Metric specs to serialize.

    Returns:
        list[dict[str, Any]]: List of metric dictionaries without None values.
    """
    return [spec.model_dump(exclude_none=True) for spec in metrics]


def _build_request_from_args_or_config(
    *,
    config: Path | None,
    input_file: Path | None,
    text_key: str | None,
    record_id_key: str | None,
    metric: list[str],
) -> EvaluationRequest:
    """Build a validated EvaluationRequest from config and CLI overrides.

    Args:
        config (Path | None): Optional config file path.
        input_file (Path | None): Optional dataset file override.
        text_key (str | None): Optional dataset text field override.
        record_id_key (str | None): Optional dataset record-id field override.
        metric (list[str]): Metric CLI values.

    Returns:
        EvaluationRequest: Validated evaluation request.

    Raises:
        typer.Exit: If merged configuration is invalid or incomplete.
    """  # noqa: DOC502
    payload: dict[str, Any] = _load_run_config(config)

    dataset = payload.setdefault("dataset", {})
    if not isinstance(dataset, dict):
        _exit("config.dataset must be an object")

    if input_file is not None:
        dataset["input_file"] = str(input_file)
    if text_key is not None:
        dataset["text_key"] = text_key
    if record_id_key is not None:
        dataset["record_id_key"] = record_id_key

    if metric:
        payload["metrics"] = _as_metric_dicts(
            _parse_metric_spec(item) for item in metric
        )

    if "metrics" not in payload:
        _exit("No metrics configured. Provide --metric or config.metrics.")

    try:
        return EvaluationRequest.create(obj=payload)
    except ValidationError as exc:
        _exit(f"Validation error: {exc}")


def _handle_evaluation_error(exc: Exception) -> NoReturn:
    """Normalize runtime exceptions into user-facing CLI errors.

    Args:
        exc (Exception): Exception raised during evaluation.

    Raises:
        typer.Exit: Always raised after printing a mapped error message.
    """  # noqa: DOC502
    msg: str
    if isinstance(
        exc,
        (
            DatasetAccessError,
            DuplicateMetricResultKeyError,
            MetricConfigurationError,
            UnsupportedDatasetError,
            UnknownMetricError,
        ),
    ):
        msg = str(exc)
    elif isinstance(exc, (ImportError, ModuleNotFoundError)):  # pragma: no cover
        msg = (
            "A required optional dependency is missing for the selected metric(s). "
            "Install the relevant extra(s) and try again."
        )
    else:  # pragma: no cover
        msg = f"An unexpected error occurred: {exc}"
    _exit(msg)


@app.command(
    "run", no_args_is_help=True, help="Evaluate a JSONL dataset on a set of metrics."
)
def run(
    config: Annotated[
        Path | None,
        typer.Option(
            "--config",
            help="JSON/YAML config with dataset and metrics.",
            rich_help_panel="File Configuration",
        ),
    ] = None,
    input_file: Annotated[
        Path | None,
        typer.Option(
            "--input-file",
            "-f",
            help="Input dataset file path.",
            rich_help_panel="Inline Configuration",
        ),
    ] = None,
    text_key: Annotated[
        str | None,
        typer.Option(
            "--text-key",
            "-t",
            help="Dataset field containing text.",
            rich_help_panel="Inline Configuration",
        ),
    ] = None,
    record_id_key: Annotated[
        str | None,
        typer.Option(
            "--record-id-key",
            "-i",
            help="Dataset field containing record IDs.",
            rich_help_panel="Inline Configuration",
        ),
    ] = None,
    metric: Annotated[
        list[str] | None,
        typer.Option(
            "--metric",
            "-m",
            help=(
                "Metric spec (repeatable). Each value may be: metric name, preset name, "
                "JSON object string, or path to JSON/YAML spec file."
            ),
            rich_help_panel="Inline Configuration",
        ),
    ] = None,
    output: Annotated[
        Path | None,
        typer.Option(
            "--output",
            "-o",
            help="Output JSONL path, defaults to stdout.",
            rich_help_panel="Inline Configuration",
        ),
    ] = None,
) -> None:
    request = _build_request_from_args_or_config(
        config=config,
        input_file=input_file,
        text_key=text_key,
        record_id_key=record_id_key,
        metric=metric or [],
    )

    try:
        stream, should_close = _open_output(output)
        try:
            count, elapsed_s = _stream_results(Evaluator(request), stream)
        finally:
            if should_close:
                stream.close()
        elapsed_time: str = (
            f"{elapsed_s // 60:.0f}m {elapsed_s % 60:.0f}s"
            if elapsed_s >= 60
            else (
                f"{elapsed_s:.2f}s"
                if 1 <= elapsed_s < 60
                else f"{elapsed_s * 1000:.2f}ms"
            )
        )
        avg_ms = (elapsed_s * 1000 / count) if count else 0.0
        err_console.rule("[bold green]Evaluation Complete[/bold green]")
        err_console.print(
            f":white_check_mark: [bold]Done[/bold] • Evaluated [bold][cyan]{count:,}[/] items[/] in [bold cyan]{elapsed_time}[/bold cyan] at an average speed of [bold][cyan]{avg_ms:.2f}ms[/]/item[/].",
            highlight=False,
            justify="center",
        )
        if output is not None and str(output) != "-":
            err_console.print(" •  •  • ", justify="center", style="dim")
            err_console.print(
                f"Saved results to '{output.name}' at [bold link={output.parent.as_uri()}]{output.parent}[/]",
                justify="center",
            )
    except typer.Exit:
        raise
    except Exception as exc:
        _handle_evaluation_error(exc)


@metrics_app.command(
    "list", help="List available metrics and their configuration arguments."
)
def metrics_list() -> None:
    def _args_renderable(
        field_items: list[tuple[str, Any]],
    ) -> str | Table:
        if not field_items:
            return "-"

        args_table = Table.grid(padding=(0, 0))
        args_table.add_column(style="bold")
        args_table.add_column(style="dim")
        args_table.add_column()
        args_table.add_column(style="dim")

        for field_name, annotation in field_items:
            args_table.add_row(
                field_name,
                " (",
                Syntax(
                    str(annotation),
                    "python",
                    word_wrap=True,
                    line_numbers=False,
                    background_color="default",
                ),
                ")",
            )
        return args_table

    metrics = sorted(MetricsCatalog.available())
    table = Table(show_header=True, header_style="bold cyan", show_lines=True)
    table.add_column("Metric Name", style="bold")
    table.add_column("Required Config Args")
    table.add_column("Optional Config Args")

    for name in metrics:
        metric_cls = MetricsCatalog.get_metric(name)
        required_args: list[tuple[str, Any]] = []
        optional_args: list[tuple[str, Any]] = []

        for field_name, field_info in metric_cls.config_schema.model_fields.items():
            if field_info.is_required():
                required_args.append((field_name, field_info.annotation))
            else:
                optional_args.append((field_name, field_info.annotation))

        required_args.sort(key=lambda item: item[0])
        optional_args.sort(key=lambda item: item[0])

        table.add_row(
            name,
            _args_renderable(required_args),
            _args_renderable(optional_args),
        )

    out_console.print(
        Panel(
            table,
            title="Available Metrics",
            border_style="green",
            expand=True,
        )
    )


@metrics_app.command(
    "info", help="Show detailed information about a registered metric."
)
def metrics_info(
    name: Annotated[
        str,
        typer.Argument(
            help="The name of the registered metric to display information for."
        ),
    ],
    format: Annotated[
        str,
        typer.Option("--format", help="Output format: text, json, yaml."),
    ] = "text",
) -> None:
    fmt = format.lower()
    if fmt not in {"text", "json", "yaml"}:  # pragma: no cover
        _exit("--format must be one of: text, json, yaml")

    try:
        metric_cls = MetricsCatalog.get_metric(name)
    except KeyError as exc:
        _exit(str(exc))

    schema = metric_cls.config_schema.model_json_schema()
    if fmt == "text":
        if not schema.get("properties"):  # pragma: no cover
            out_console.print(
                Panel(
                    Group(
                        Text(
                            f"config_model: {metric_cls.config_schema.__name__}",
                            style="cyan",
                        ),
                        Text("Runtime config: none"),
                    ),
                    title=Text(f"Metric: {name}", style="bold"),
                    border_style="green",
                )
            )
            return
        field_panels: list[Panel] = []
        for field_name, field_info in metric_cls.config_schema.model_fields.items():
            meta_table = Table.grid(padding=(0, 2))
            meta_table.add_column(style="bold")
            meta_table.add_column()
            meta_table.add_row("required", Text(str(field_info.is_required())))
            meta_table.add_row(
                "annotation",
                Syntax(
                    str(field_info.annotation),
                    "python",
                    word_wrap=True,
                    line_numbers=False,
                    background_color="default",
                ),
            )
            if not field_info.is_required():  # pragma: no cover
                meta_table.add_row("default", Text(repr(field_info.default)))

            description_panel = Panel(
                Text(
                    field_info.description or "No description provided.",
                    overflow="fold",
                ),
                title="description",
            )

            field_panels.append(
                Panel(
                    Group(meta_table, description_panel),
                    title=field_name,
                    border_style="cyan",
                )
            )

        out_console.print(
            Panel(
                Group(
                    Text(
                        f"config_model: {metric_cls.config_schema.__name__}",
                        style="cyan",
                    ),
                    Group(*field_panels),
                ),
                title=Text(f"Metric: {name}", style="bold"),
                style="bold",
                border_style="green",
                expand=True,
            )
        )
        return

    if fmt == "json":
        typer.echo(orjson.dumps(schema, option=orjson.OPT_INDENT_2).decode())
        return

    typer.echo(yaml.safe_dump(schema, sort_keys=False))


@metrics_app.command("presets", help="Show available configuration presets.")
def metrics_presets() -> None:
    table = Table(show_header=True, header_style="bold cyan")
    table.add_column("Preset", style="bold")
    table.add_column("Purpose")

    table.add_row("count_characters", "Count characters in evaluated text")
    table.add_row("count_tokens", "Count tokens")
    table.add_row("count_tokens:wordpunct", "Count tokens with wordpunct segmentation")
    table.add_row(
        "count_tokens:whitespace",
        "Count tokens with whitespace segmentation (same as [bold]count_tokens[/bold])",
    )
    table.add_row(
        "count_encodings:<hf-model-name>",
        "Count model-specific encodings using a HuggingFace tokenizer",
    )

    out_console.print(
        Panel(
            table,
            title="Convenience Presets",
            border_style="green",
            expand=True,
        )
    )


if __name__ == "__main__":  # pragma: no cover
    app()
