# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/Benchmark_CLI.ipynb.

# %% auto 0
__all__ = ['BENCHMARK_RESULTS_DIR', 'app', 'benchmark']

# %% ../../nbs/Benchmark_CLI.ipynb 1
from typing import *
import shutil
from tempfile import TemporaryDirectory
from pathlib import Path
from contextlib import contextmanager
import os
import subprocess # nosec: B404: Consider possible security implications associated with the subprocess module.

import typer
from yaspin import yaspin

import numpy as np

from .._code_generator.constants import INCOMPLETE_APP_ERROR_MSG


# %% ../../nbs/Benchmark_CLI.ipynb 3
BENCHMARK_RESULTS_DIR = "benchmark-results"

# %% ../../nbs/Benchmark_CLI.ipynb 4
@contextmanager
def _set_cwd(cwd_path: Union[Path, str]) -> Generator:
    """Set the current working directory for the duration of the context manager.

    Args:
        cwd_path: The path to the new working directory.

    !!! note

        The above docstring is autogenerated by docstring-gen library (https://github.com/airtai/docstring-gen)
    """
    cwd_path = Path(cwd_path)
    original_cwd = os.getcwd()
    os.chdir(cwd_path)

    try:
        yield
    finally:
        os.chdir(original_cwd)

# %% ../../nbs/Benchmark_CLI.ipynb 5
app = typer.Typer(
    short_help="Run benchmark against pre-defined example app descriptions",
)

# %% ../../nbs/Benchmark_CLI.ipynb 6
@app.command(
    "run_benchmark",
    help="Run benchmark against pre-defined example app descriptions",
)
def benchmark(
    fixtures_path: str = typer.Argument(
        ...,
        help="The path to the pre-defined example app descriptions",
    ),
    no_repeat: int = typer.Option(
        1,
        "--repeat",
        "-r",
        help="Number of generation repetitions per app description",
    ),
) -> None:
    fixtures_path_obj = Path(fixtures_path).resolve()

    app_descriptions = [
        filename
        for filename in fixtures_path_obj.glob("*.txt")
        if "-log" not in filename.stem
    ]
    
    if no_repeat > 1:
        app_descriptions = np.repeat(np.array(app_descriptions), no_repeat).tolist()
    
    no_of_description_files = len(app_descriptions)
    typer.secho(
        f"Total app description files: {no_of_description_files}", fg=typer.colors.CYAN
    )
    success_cnt = 0
    for i, app_description in enumerate(app_descriptions):
        with yaspin(
            text=f"{i+1}/{no_of_description_files} Generating app for: {app_description.name}",
            color="cyan",
            spinner="clock",
        ) as sp:
            with TemporaryDirectory() as d:
                with _set_cwd(d):
                    cli_command = (
                        f"faststream_gen -i {app_description} -o {d}/new-project --dev"
                    )

                    # nosemgrep: python.lang.security.audit.subprocess-shell-true.subprocess-shell-true
                    result = subprocess.run(  # nosec: B602, B603 subprocess call - check for execution of untrusted input.
                        cli_command,
                        shell=True,
                        check=True,
                        stdout=subprocess.PIPE,
                        stderr=subprocess.PIPE,
                        text=True,
                    )
                    sp.text = ""
                    if INCOMPLETE_APP_ERROR_MSG in str(result.stdout):
                        sp.text = ""
                        sp.color = "red"
                        sp.ok(
                            f" ✘ Error: App generated failed for: {app_description.name}"
                        )

                    else:
                        sp.ok(
                            f" ✔ App successfully generated for: {app_description.name}"
                        )
                        success_cnt += 1

                    benchmark_results_dir = fixtures_path_obj / BENCHMARK_RESULTS_DIR
                    benchmark_results_dir.mkdir(parents=True, exist_ok=True)
                    shutil.copytree(
                        str(Path(d) / "new-project"),
                        benchmark_results_dir / Path(app_description).stem,
                        dirs_exist_ok=True,
                    )

    success_rate = success_cnt / no_of_description_files
    typer.secho(
        f"Success rate: {success_rate}",
        fg=typer.colors.CYAN,
    )
    if int(success_rate) != 1:
        typer.secho(
            f"\nTo debug, go to {str(fixtures_path_obj / BENCHMARK_RESULTS_DIR)}, select one of the failed app description directories, and run 'pytest'.",
            fg=typer.colors.RED,
        )
