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_response(user_input: str) -> str:
#     template = """
#     You are an intelligent chatbot.
#     Generate a short response answering the user's input: {user_input}. 
#     and ask at the end "what can I do for you now"
#     """
#     prompt = ChatPromptTemplate.from_template(template)
#     chain = prompt | llm | StrOutputParser()
#     response = chain.invoke({"user_input": user_input})
#     return response    


# Function remains unchanged, but now supports character context
def generate_response(
    user_input: str,
    name: str = "Assistant",  # Default name if not provided
    gender: str = "neutral",  # Default gender if not provided
    backstory: str = "A helpful virtual assistant."  # Default backstory if not provided
) -> str:
    template = """
    You are an intelligent chatbot.
    Generate a short response answering the user's input: {user_input}.
    and ask for other queries at the end.
    """
    # Combine character context with the original prompt
    full_prompt = RunnableParallel({
        "character_context": ChatPromptTemplate.from_template(CHARACTER_TEMPLATE),
        "task_instructions": ChatPromptTemplate.from_template(template)
    })

    # Final assembly template
    final_template = """
    {character_context}

    {task_instructions}

    Additional Guidelines:
    - Use your character's unique voice and style in the response.
    - Ensure the response aligns with your backstory.
    - Keep the response natural and conversational.
    """

    # Update the chain to include character context
    chain = (
        full_prompt |
        ChatPromptTemplate.from_template(final_template) |
        llm |
        StrOutputParser()
    )
    """Generate a response to the user's input, incorporating character context."""
    response = chain.invoke({
        "name": name,
        "gender": gender,
        "backstory": backstory,
        "user_input": user_input,
    })
    return response