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

# %% auto 0
__all__ = ['create_venv_and_run_tests_bash_script', 'fix_requirements_and_run_tests']

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

from yaspin import yaspin

from .logger import get_logger

from .._code_generator.prompts import REQUIREMENTS_GENERATION_PROMPT
from faststream_gen._code_generator.constants import (
    APPLICATION_FILE_PATH,
    TEST_FILE_PATH,
    STEP_LOG_DIR_NAMES,
    TOML_FILE_NAME,
    OpenAIModel,
)

from .._code_generator.chat import CustomAIChat, ValidateAndFixResponse

from faststream_gen._code_generator.helper import (
    download_and_extract_faststream_archive,
    write_file_contents,
    read_file_contents,
    set_cwd,
    mock_openai_create,
    retry_on_error,
)

# %% ../../nbs/Integration_Test_Generator.ipynb 3
create_venv_and_run_tests_bash_script = """
#!/bin/bash

# get the project directory from command line argument
project_dir=$1

# get the venv directory from command line argument
venv_dir=$2

# create a new venv in the specified directory
python3 -m venv $venv_dir/my_venv > /dev/null 2>&1

# activate the venv
source $venv_dir/my_venv/bin/activate

# navigate to the project directory
cd $project_dir

# install the python project inside the venv
pip install .['dev'] > /dev/null 2>&1

# run pytest and capture output
pytest_output=$(pytest --tb=short)

# capture pytest exit code
pytest_exit_code=$?

# print the pytest output
echo "pytest_output_start:$pytest_output:pytest_output_end"

# print the pytest exit code
echo "pytest_exit_code:$pytest_exit_code"

# deactivate the venv
deactivate
"""

# %% ../../nbs/Integration_Test_Generator.ipynb 4
def _setup_venv_and_run_tests(output_path: str) -> List[str]:
    output_path_resolved = Path(output_path).resolve()
    with TemporaryDirectory() as d:
        bash_file = Path(d) / "run_tests.sh"
        write_file_contents(str(bash_file), create_venv_and_run_tests_bash_script)
        with set_cwd(d):
            # nosemgrep: python.lang.security.audit.subprocess-shell-true.subprocess-shell-true
            p = subprocess.run( # nosec: B602, B603, B607 subprocess call - check for execution of untrusted input.
                ["bash", "run_tests.sh", output_path_resolved, Path(d).resolve()],
                capture_output=True,
                text=True,
            )
            
        # Extract exit code
        exit_code = int(re.search('pytest_exit_code:(\d+)', p.stdout).group(1)) # type: ignore
        if exit_code !=0:
            # Extract pytest output
            pytest_output: str = re.search('pytest_output_start:(.*):pytest_output_end', p.stdout, re.DOTALL).group(1).strip() # type: ignore
            return [pytest_output]
        
        return []    

# %% ../../nbs/Integration_Test_Generator.ipynb 7
def _stript(s: str) -> str:
    return s.strip().strip('"')


def _split_app_and_test_req(response: str) -> Tuple[str, str]:

    app_req, test_req = response.split("==== APP REQUIREMENT ====")[1].split(
        "==== TEST REQUIREMENT ===="
    )
    return _stript(app_req), _stript(test_req)

# %% ../../nbs/Integration_Test_Generator.ipynb 10
def _update_toml_file(output_dir: str, app_req: str, test_req: str) -> None:
    toml_file_path = f"{output_dir}/{TOML_FILE_NAME}"
    toml_contents = read_file_contents(toml_file_path)
    data = toml.loads(toml_contents)
    
    app_reqs = [r.strip() for r in app_req.split(",")]
    test_reqs = [r.strip() for r in test_req.split(",")]
    test_reqs = [r for r in test_reqs if r != "pytest"]
    
    data["project"]["dependencies"] = data["project"]["dependencies"] + app_reqs
    data["project"]["optional-dependencies"]["testing"] = data["project"]["optional-dependencies"]["testing"] + test_reqs
    
    toml_string = toml.dumps(data)
    write_file_contents(toml_file_path, toml_string)

# %% ../../nbs/Integration_Test_Generator.ipynb 12
def _validate_response(
    response: str, output_directory: str, **kwargs: Dict[str, Any]
) -> List[str]:
    try:
        app_req, test_req = _split_app_and_test_req(response)
    except (IndexError, ValueError) as e:
        return [
            "Please add ==== APP REQUIREMENT ==== and ==== TEST REQUIREMENT ==== in your response"
        ]
    
    _update_toml_file(output_directory, app_req, test_req)
    
    return _setup_venv_and_run_tests(output_directory)

# %% ../../nbs/Integration_Test_Generator.ipynb 16
@retry_on_error()  # type: ignore
def _generate(
    model: str,
    prompt: str,
    app_and_test_code: str,
    total_usage: List[Dict[str, int]],
    output_directory: str,
    **kwargs,
) -> Tuple[str, List[Dict[str, int]], bool]:
    requirements_generator = CustomAIChat(
        params={
            "temperature": 0.2,
        },
        model=model,
        user_prompt=prompt,
    )
    requirements_validator = ValidateAndFixResponse(
        requirements_generator, _validate_response
    )
    validator_result = requirements_validator.fix(
        app_and_test_code,
        total_usage,
        STEP_LOG_DIR_NAMES["requirements"],
        str(output_directory),
        **kwargs,
    )
    return (
        (validator_result, True) # type: ignore
        if isinstance(validator_result[-1], defaultdict)
        else validator_result
    )

# %% ../../nbs/Integration_Test_Generator.ipynb 19
def fix_requirements_and_run_tests(
    output_directory: str,
    model: str,
    total_usage: List[Dict[str, int]],
) -> Tuple[List[Dict[str, int]], bool]:
    with yaspin(
        text="Running integration tests...", color="cyan", spinner="clock"
    ) as sp:
        app_file = Path(output_directory) / APPLICATION_FILE_PATH
        app_code = read_file_contents(str(app_file))

        test_file = Path(output_directory) / TEST_FILE_PATH
        test_code = read_file_contents(str(test_file))

        app_and_test_code = f"==== APP CODE ====\n\n{app_code}\n\n==== TEST CODE ====\n\n{test_code}\n\n"

        total_usage, is_requirements_file_valid = _generate(
            model,
            REQUIREMENTS_GENERATION_PROMPT,
            app_and_test_code,
            total_usage,
            output_directory,
        )

        sp.text = ""
        if is_requirements_file_valid:
            message = " ✔ Integration tests were successfully completed."
        else:
            message = " ✘ Error: Integration tests failed."
            sp.color = "red"

        sp.ok(message)

        return total_usage, is_requirements_file_valid
