from __future__ import annotations

import json
from pathlib import Path
from typing import Optional

import anyio
import typer
from rich.console import Console

from leshy.cli.app import app
from leshy.core.config import load_config
from leshy.core.models import TestSuiteResult
from leshy.core.runner import Runner
from leshy.generators.guardrail import GuardrailGenerator

console = Console()

_DEFAULT_OUTPUT_DIR = Path("leshy") / "runs"


def _run_tests(
    config_path: Path,
    tag: Optional[str],
    verbose: bool,
) -> TestSuiteResult:
    """Load config and run the test suite. Returns TestSuiteResult."""
    try:
        config = load_config(config_path)
    except (FileNotFoundError, ValueError) as e:
        typer.echo(f"Error loading config: {e}", err=True)
        raise typer.Exit(code=2)

    runner = Runner(config, tag_filter=tag, verbose=verbose)

    try:
        result = anyio.from_thread.run(runner.run)
    except Exception:
        result = anyio.run(runner.run)

    return result


def _save_results_json(result: TestSuiteResult, path: str) -> None:
    """Serialize test results to JSON file."""
    out = Path(path)
    out.parent.mkdir(parents=True, exist_ok=True)
    data = result.model_dump(mode="json")
    out.write_text(json.dumps(data, indent=2, default=str))
    console.print(f"[dim]Results saved to {out}[/dim]")


def _generate_guardrails(
    generator: GuardrailGenerator,
    output: Optional[str],
) -> None:
    """Generate guardrail node and write to file."""
    if output:
        output_path = Path(output)
    else:
        output_path = _DEFAULT_OUTPUT_DIR / "guardrail_node.py"

    output_path.parent.mkdir(parents=True, exist_ok=True)

    if output_path.exists():
        overwrite = typer.confirm(f"{output_path} already exists. Overwrite?")
        if not overwrite:
            raise typer.Abort()

    result_path = generator.write(output_path)

    vuln_categories = generator._find_vulnerable_categories()
    if not vuln_categories:
        console.print(
            "\n[yellow]No vulnerabilities detected.[/yellow] "
            "No guardrails generated.\n"
        )
        return

    console.print(
        f"\n[green bold]Guardrail nodes generated:[/green bold] {result_path}"
    )
    console.print(
        f"[dim]Categories covered: {', '.join(sorted(vuln_categories.keys()))}[/dim]"
    )
    console.print(
        "[dim]Import the guardrail functions into your LangGraph workflow.[/dim]\n"
    )


def _print_summary(result: TestSuiteResult) -> None:
    """Print a short test results summary."""
    if result.total_vulnerable > 0:
        console.print(
            f"\n[red bold]Vulnerabilities found:[/red bold] "
            f"{result.total_vulnerable}/{result.total_payloads} payloads"
        )
    else:
        console.print(
            f"\n[green bold]All clean:[/green bold] "
            f"{result.total_resilient}/{result.total_payloads} resilient"
        )


@app.command()
def run(
    config_path: Path = typer.Argument(
        "leshy_config.yaml", help="Path to YAML config file"
    ),
    test: bool = typer.Option(
        False, "--test", help="Run security tests only"
    ),
    guardrails: bool = typer.Option(
        False, "--guardrails", help="Run tests + generate guardrail nodes"
    ),
    from_results: Optional[Path] = typer.Option(
        None, "--from-results", help="Generate guardrails from existing JSON results"
    ),
    tag: Optional[str] = typer.Option(
        None, "--tag", "-t", help="Only run tests with this tag"
    ),
    verbose: bool = typer.Option(
        False, "--verbose", help="Show per-payload results"
    ),
    output: Optional[str] = typer.Option(
        None, "--output", "-o", help="Guardrail output file path"
    ),
    save_results: Optional[str] = typer.Option(
        None, "--save-results", "-r", help="Save test results to JSON"
    ),
) -> None:
    """Run security tests and generate guardrail nodes."""

    # Mode 2: generate from existing results file
    if from_results is not None:
        if not from_results.exists():
            typer.echo(f"Results file not found: {from_results}", err=True)
            raise typer.Exit(code=2)
        try:
            generator = GuardrailGenerator(from_results)
            _generate_guardrails(generator, output)
        except Exception as e:
            typer.echo(f"Error generating guardrails: {e}", err=True)
            raise typer.Exit(code=2)
        return

    if not config_path.exists():
        typer.echo(f"Config file not found: {config_path}", err=True)
        raise typer.Exit(code=2)

    # Determine mode
    if not test and not guardrails:
        # Interactive prompt
        console.print("\n[bold]What would you like to do?[/bold]")
        console.print("  [cyan][1][/cyan] Run security tests")
        console.print("  [cyan][2][/cyan] Run tests + generate guardrail nodes")
        choice = typer.prompt("Choose", default="1", show_default=True)
        if choice == "2":
            guardrails = True
        else:
            test = True

    # Mode 1: test only
    if test:
        result = _run_tests(config_path, tag, verbose)
        _print_summary(result)
        if save_results:
            _save_results_json(result, save_results)
        if result.total_vulnerable > 0:
            raise typer.Exit(code=1)
        if result.total_errors > 0 and result.total_errors == result.total_payloads:
            raise typer.Exit(code=2)
        return

    # Mode 3: test + generate guardrails
    if guardrails:
        result = _run_tests(config_path, tag, verbose)
        _print_summary(result)
        if save_results:
            _save_results_json(result, save_results)

        try:
            generator = GuardrailGenerator.from_result(result)
            _generate_guardrails(generator, output)
        except Exception as e:
            typer.echo(f"Error generating guardrails: {e}", err=True)
            raise typer.Exit(code=2)

        if result.total_vulnerable > 0:
            raise typer.Exit(code=1)
        if result.total_errors > 0 and result.total_errors == result.total_payloads:
            raise typer.Exit(code=2)
