from langchain.prompts import ChatPromptTemplate # type: ignore
from langchain_core.output_parsers import StrOutputParser # type: ignore
from Ai_Agents.models.language_model import llm
from langchain_core.runnables import RunnableParallel # type: ignore
from Ai_Agents.services.templates.character_template import CHARACTER_TEMPLATE


def generate_general_response(
    user_input: str,
    objective: str,
    name: str,
    gender: str,
    backstory: str
) -> str:
    # Original template remains unchanged
    base_template = """
    Your objective: {objective}
    The user asked an unrelated question. Respond while:
    1. Briefly answering their query: {user_input}
    2. Maintaining subtle alignment with the objective
    3. Keeping responses under 3 sentences
    """
    
    # Chain character context with original template
    full_prompt = RunnableParallel({
        "character_context": ChatPromptTemplate.from_template(CHARACTER_TEMPLATE),
        "task_instructions": ChatPromptTemplate.from_template(base_template)
    })
    
    # Combined final template
    final_template = """
    {character_context}
    
    {task_instructions}
    
    Finalize your response using:
    - Your character's communication style
    - The required objective alignment
    - Natural conversation flow
    """
    
    chain = (
        full_prompt | 
        ChatPromptTemplate.from_template(final_template) | 
        llm | 
        StrOutputParser()
    )
    
    return chain.invoke({
        "name": name,
        "gender": gender,
        "backstory": backstory,
        "objective": objective,
        "user_input": user_input
    })