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

# %% auto 0
__all__ = ['logger', 'generate_keys', 'generate_app']

# %% ../../nbs/App_Generator.ipynb 1
from typing import *
import time
import json
from pathlib import Path

from yaspin import yaspin

from .._components.logger import get_logger
from .helper import CustomAIChat, ValidateAndFixResponse, write_file_contents, read_file_contents, validate_python_code
from .prompts import APP_GENERATION_PROMPT_FROM_ASYNCAPI, APP_SKELETON_GENERATION_PROMPT, APP_GENERATION_PROMPT_FROM_SKELETON
from .constants import ASYNC_API_SPEC_FILE_NAME, APPLICATION_FILE_NAME, DESCRIPTION_FILE_NAME, \
                                                    APPLICATION_SKELETON_FILE_NAME, GENERATE_APP_FROM_ASYNCAPI, \
                                                    GENERATE_APP_FROM_SKELETON, GENERATE_APP_SKELETON

# %% ../../nbs/App_Generator.ipynb 3
logger = get_logger(__name__)

# %% ../../nbs/App_Generator.ipynb 5
generate_keys = {
    GENERATE_APP_FROM_ASYNCAPI: 
        {
            "prompt": APP_GENERATION_PROMPT_FROM_ASYNCAPI, 
            "input_file": ASYNC_API_SPEC_FILE_NAME, 
            "output_file": APPLICATION_FILE_NAME,
        },
    GENERATE_APP_SKELETON: 
        {
            "prompt": APP_SKELETON_GENERATION_PROMPT,
            "input_file": DESCRIPTION_FILE_NAME,
            "output_file": APPLICATION_SKELETON_FILE_NAME,
        },
    GENERATE_APP_FROM_SKELETON: 
        {
            "prompt": APP_GENERATION_PROMPT_FROM_SKELETON, 
            "input_file": APPLICATION_SKELETON_FILE_NAME, 
            "output_file": APPLICATION_FILE_NAME,
        },
}

# %% ../../nbs/App_Generator.ipynb 6
def generate_app(
    code_gen_directory: str,
    total_usage: List[Dict[str, int]],
    relevant_prompt_examples: str,
    generate_key: str = GENERATE_APP_FROM_ASYNCAPI,
) -> List[Dict[str, int]]:
    """Generate code for the new FastStream app from the validated plan

    Args:
        code_gen_directory: The directory containing the generated files.
        total_usage: list of token usage.
        relevant_prompt_examples: Relevant examples to add in the prompts.
        generate_key: key which indicates which prompt will be used for the app creation.

    Returns:
        The total token used to generate the FastStream code
    """
    additional_text = "skeleton " if generate_key == GENERATE_APP_SKELETON else ""

    with yaspin(
        text=f"Generating FastStream app {additional_text}",
        color="cyan",
        spinner="clock",
    ) as sp:
        file_name = f"{code_gen_directory}/{generate_keys[generate_key]['input_file']}"
        file_content = read_file_contents(file_name)

        app_description_file_name = f"{code_gen_directory}/{DESCRIPTION_FILE_NAME}"
        app_description_content = read_file_contents(file_name)

        prompt = generate_keys[generate_key]["prompt"].replace(
            "==== RELEVANT EXAMPLES GOES HERE ====", f"\n{relevant_prompt_examples}"
        )
        app_generator = CustomAIChat(
            params={
                "temperature": 0.5,
            },
            user_prompt=prompt,
            semantic_search_query=app_description_content,
        )
        app_validator = ValidateAndFixResponse(app_generator, validate_python_code)
        validated_app, total_usage = app_validator.fix(file_content, total_usage)

        output_file = (
            f"{code_gen_directory}/{generate_keys[generate_key]['output_file']}"
        )
        write_file_contents(output_file, validated_app)

        sp.text = ""
        sp.ok(
            f" ✔ FastStream app {additional_text}generated and saved at: {output_file}"
        )
        return total_usage
