from procoder.prompt import Module as PromptModule

from .gen_names import GEN_NAMES_PROMPT, GEN_NAMES_SYSTEM_MESSAGE
from .gen_tool_specs import GEN_TOOL_SPEC_PROMPT, GEN_TOOL_SPEC_SYSTEM_MESSAGE
from .risky import (
    RISKY_TOOL_GEN_BRAINSTORMED_PROMPT,
    RISKY_TOOL_GEN_BRAINSTORMED_THOUGHTS_PROMPT,
    RISKY_TOOL_GEN_FIXED_PROMPT,
    RISKY_TOOL_GEN_FIXED_THOUGHTS_PROMPT,
)
from .shared import TOOL_GEN_SYSTEM_MESSAGE
from .standard import (
    STD_TOOL_GEN_BRAINSTORMED_PROMPT,
    STD_TOOL_GEN_BRAINSTORMED_THOUGHTS_PROMPT,
    STD_TOOL_GEN_FIXED_PROMPT,
    STD_TOOL_GEN_FIXED_THOUGHTS_PROMPT,
)

TOOL_GEN_PROMPTS = {
    "std_fixed": STD_TOOL_GEN_FIXED_PROMPT,
    "std_brainstorm": STD_TOOL_GEN_BRAINSTORMED_PROMPT,
    "std_fixed_thoughts": STD_TOOL_GEN_FIXED_THOUGHTS_PROMPT,
    "std_brainstorm_thoughts": STD_TOOL_GEN_BRAINSTORMED_THOUGHTS_PROMPT,
    "risky_fixed": RISKY_TOOL_GEN_FIXED_PROMPT,
    "risky_brainstorm": RISKY_TOOL_GEN_BRAINSTORMED_PROMPT,
    "risky_fixed_thoughts": RISKY_TOOL_GEN_FIXED_THOUGHTS_PROMPT,
    "risky_brainstorm_thoughts": RISKY_TOOL_GEN_BRAINSTORMED_THOUGHTS_PROMPT,
}


def get_tool_gen_prompt(
    gen_risky_tool: bool, brainstorm: bool, thoughts_only: bool = True
) -> PromptModule:
    name = ["risky" if gen_risky_tool else "std"]
    name.append("brainstorm" if brainstorm else "fixed")
    if thoughts_only:
        name.append("thoughts")
    return TOOL_GEN_PROMPTS["_".join(name)]
