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

# %% auto 0
__all__ = ['logger', 'CODE_CONTAINS_IMPLEMENTATION_ERROR', 'generate_app_skeleton']

# %% ../../nbs/App_Skeleton_Generator.ipynb 1
from typing import *
import time
import json
import ast
from pathlib import Path
from collections import defaultdict

from yaspin import yaspin

from .._components.logger import get_logger
from .chat import CustomAIChat, ValidateAndFixResponse
from faststream_gen._code_generator.helper import (
    write_file_contents,
    validate_python_code,
    retry_on_error,
)
from .prompts import APP_SKELETON_GENERATION_PROMPT
from faststream_gen._code_generator.constants import (
    STEP_LOG_DIR_NAMES,
    APPLICATION_FILE_PATH,
    LOGS_DIR_NAME
)

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

# %% ../../nbs/App_Skeleton_Generator.ipynb 5
CODE_CONTAINS_IMPLEMENTATION_ERROR = "Error: The response contains code implementation. Rewrite the skeleton code without implementing the business logic for the functions. Ensure the new code has only google styled docstring describing the business logic step by step and raise NotImplementedError()"

# %% ../../nbs/App_Skeleton_Generator.ipynb 6
def _has_function_implementation(
    node: Union[ast.AsyncFunctionDef, ast.FunctionDef]
) -> bool:
    return len(node.body) == 2 and isinstance(node.body[-1], ast.Raise)


def _check_response_for_implementation(response: str) -> List[str]:
    parsed = ast.parse(response)
    function_nodes = [
        node
        for node in ast.walk(parsed)
        if isinstance(node, ast.AsyncFunctionDef) or isinstance(node, ast.FunctionDef)
    ]
    ret_val = (
        []
        if all(_has_function_implementation(node) for node in function_nodes)
        else [CODE_CONTAINS_IMPLEMENTATION_ERROR]
    )
    return ret_val

# %% ../../nbs/App_Skeleton_Generator.ipynb 9
def _validate_response(
    response: str, output_directory: str, **kwargs: Dict[str, Any]
) -> List[str]:
    target_file_name = Path(output_directory) / APPLICATION_FILE_PATH
    write_file_contents(str(target_file_name), response)
    code_import_errors = validate_python_code(str(target_file_name))
    
    if len(code_import_errors) != 0:
        return code_import_errors
    
    return _check_response_for_implementation(response)

# %% ../../nbs/App_Skeleton_Generator.ipynb 15
@retry_on_error()  # type: ignore
def _generate(
    model: str,
    prompt: str,
    app_description_content: str,
    total_usage: List[Dict[str, int]],
    output_directory: str,
    **kwargs,
) -> Tuple[str, List[Dict[str, int]]]:
    app_generator = CustomAIChat(
        params={
            "temperature": 0.2,
        },
        model=model,
        user_prompt=prompt,
        #         semantic_search_query=app_description_content,
    )
    app_validator = ValidateAndFixResponse(app_generator, _validate_response)
    validator_result = app_validator.fix(
        app_description_content,
        total_usage,
        STEP_LOG_DIR_NAMES["skeleton"],
        str(output_directory),
        **kwargs,
    )

    return (
        (validator_result, True)
        if isinstance(validator_result[-1], defaultdict)
        else validator_result
    )

# %% ../../nbs/App_Skeleton_Generator.ipynb 18
def generate_app_skeleton(
    validated_description: str,
    output_directory: str,
    model: str,
    total_usage: List[Dict[str, int]],
    relevant_prompt_examples: str,
) -> Tuple[List[Dict[str, int]], bool]:
    """Generate skeleton code for the new FastStream app from the application description

    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.

    Returns:
        The total token used to generate the FastStream code
    """
    logger.info("==== Description to Skeleton Generation ====")
    with yaspin(
        text=f"Generating FastStream app skeleton code (usually takes around 15 to 45 seconds)...",
        color="cyan",
        spinner="clock",
    ) as sp:
        prompt = APP_SKELETON_GENERATION_PROMPT.replace(
            "==== RELEVANT EXAMPLES GOES HERE ====", f"\n{relevant_prompt_examples}"
        )

        total_usage, is_valid_skeleton_code = _generate(
            model, prompt, validated_description, total_usage, output_directory
        )
        sp.text = ""
        if is_valid_skeleton_code:
            message = f" ✔ FastStream app skeleton code generated."
        else:
            message = " ✘ Error: Failed to generate a valid application skeleton code."
            sp.color = "red"

        sp.ok(message)
        return total_usage, is_valid_skeleton_code
