from __future__ import annotations

from pathlib import Path

import typer
from rich.console import Console

from leshy.cli.app import app
from leshy.cli.wizard import (
    generate_custom_test_yaml,
    generate_main_config_yaml,
    generate_test_yaml_for_category,
    run_wizard,
)

console = Console()

_STARTER_YAML = """\
# Leshy security test configuration
# Docs: https://github.com/bardsai/leshy

name: My Security Scan

# LLM model for judge evaluation and LLM buffs (powered by litellm)
# Supports any litellm model string:
#   OpenAI:    gpt-4o-mini, gpt-4o          (set OPENAI_API_KEY)
#   Anthropic: claude-sonnet-4-5-20250929          (set ANTHROPIC_API_KEY)
#   Google:    gemini/gemini-1.5-flash       (set GEMINI_API_KEY)
#   Azure:     azure/gpt-4o-mini             (set AZURE_API_KEY + AZURE_API_BASE)
llm_model: gpt-4o-mini
# api_key: your-key-here  # or omit to use provider env var (e.g. OPENAI_API_KEY)

# Number of payloads to send in parallel (per test case)
concurrency: 5

target:
  type: http
  url: https://my-agent.example.com/chat
  headers:
    Authorization: "Bearer ${API_KEY}"
  body_template: '{"prompt": "{{payload}}"}'  # {{payload}} = attack text
  # For multi-turn tests, add session tracking:
  # body_template: '{"prompt": "{{payload}}", "session_id": "{{session_id}}"}'
  response_path: "$.response"

loggers:
  - console
  - json_file

tests_dir: leshy/templates

"""


@app.command()
def init(
    output: Path = typer.Option(
        "leshy_config.yaml", "--output", "-o", help="Output file path"
    ),
    quick: bool = typer.Option(
        False, "--quick", "-q", help="Skip wizard, write starter template"
    ),
) -> None:
    """Generate a Leshy config file — interactively or from a starter template."""
    if output.exists():
        overwrite = typer.confirm(f"{output} already exists. Overwrite?")
        if not overwrite:
            raise typer.Abort()

    if quick:
        output.write_text(_STARTER_YAML)
        _create_test_directory(["prompt_injection", "jailbreak", "data_extraction"])
        typer.echo(f"Config written to {output}")
        return

    # Interactive wizard
    result = run_wizard()

    # Generate main config
    config_yaml = generate_main_config_yaml(result)
    output.write_text(config_yaml)

    console.print(f"\n[green bold]Config written to {output}[/green bold]")

    # Always generate leshy/ folder with test files
    _create_test_directory(result.selected_categories)

    # Print warnings and next steps
    if result.env_warnings and not result.api_key:
        console.print(
            "\n[yellow]Set these environment variables before running tests:[/yellow]"
        )
        for var in result.env_warnings:
            console.print(f"  export {var}=your-key-here")

    console.print(f"\n[dim]You can edit {output} to adjust settings anytime.[/dim]")
    console.print(f"[dim]Next: leshy run {output}[/dim]\n")


def _create_test_directory(categories: list[str]) -> None:
    """Create leshy/ directory with templates/ for tests and runs/ for outputs."""
    base_dir = Path("leshy")
    templates_dir = base_dir / "templates"

    for cat in categories:
        cat_dir = templates_dir / cat
        cat_dir.mkdir(parents=True, exist_ok=True)
        (cat_dir / "tests.yaml").write_text(generate_test_yaml_for_category(cat))

    custom_dir = templates_dir / "custom"
    custom_dir.mkdir(parents=True, exist_ok=True)
    (custom_dir / "tests.yaml").write_text(generate_custom_test_yaml())

    runs_dir = base_dir / "runs"
    runs_dir.mkdir(parents=True, exist_ok=True)

    console.print(f"[dim]Test directory created at {base_dir}/[/dim]")
