from __future__ import annotations

from pathlib import Path
from typing import Optional

import typer
from rich.console import Console

from leshy.cli.app import app
from leshy.cli.wizard import (
    ALL_CATEGORIES,
    CATEGORY_LABELS,
    generate_custom_test_yaml,
    generate_test_yaml_for_category,
)

console = Console()


@app.command()
def generate(
    output_dir: Path = typer.Option(
        "leshy/templates", "--output-dir", "-o", help="Output directory for test files"
    ),
    categories: Optional[str] = typer.Option(
        None, "--categories", "-c",
        help="Comma-separated categories (e.g. prompt_injection,jailbreak). Default: all",
    ),
) -> None:
    """Generate test YAML files organized by attack category."""

    if categories:
        selected = [c.strip() for c in categories.split(",")]
        invalid = [c for c in selected if c not in ALL_CATEGORIES]
        if invalid:
            typer.echo(
                f"Unknown categories: {', '.join(invalid)}. "
                f"Available: {', '.join(ALL_CATEGORIES)}",
                err=True,
            )
            raise typer.Exit(code=2)
    else:
        selected = list(ALL_CATEGORIES)

    if output_dir.exists() and any(output_dir.iterdir()):
        overwrite = typer.confirm(f"{output_dir}/ already has files. Continue?")
        if not overwrite:
            raise typer.Abort()

    created: list[str] = []

    for cat in selected:
        cat_dir = output_dir / cat
        cat_dir.mkdir(parents=True, exist_ok=True)
        test_file = cat_dir / "tests.yaml"
        test_file.write_text(generate_test_yaml_for_category(cat))
        label = CATEGORY_LABELS.get(cat, cat)
        created.append(f"  {test_file} ({label})")

    # Custom template
    custom_dir = output_dir / "custom"
    custom_dir.mkdir(parents=True, exist_ok=True)
    custom_file = custom_dir / "tests.yaml"
    if not custom_file.exists():
        custom_file.write_text(generate_custom_test_yaml())
        created.append(f"  {custom_file} (Custom)")

    console.print("\n[green bold]Test files generated:[/green bold]")
    for line in created:
        console.print(f"[dim]{line}[/dim]")

    console.print(
        f"\n[dim]Add [bold]tests_dir: {output_dir}[/bold] to your config"
        " to use these tests.[/dim]\n"
    )
