from .shared_state import SharedState


from langchain_google_genai import GoogleGenerativeAI

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import BaseMessage, SystemMessage

from typing import Dict, List


class Assistant:
    def __init__(self):
        self._llm = GoogleGenerativeAI(model="gemini-2.0-flash")
        self._prompt_template = ChatPromptTemplate.from_messages([
            SystemMessage(
                content="""You are a helpful assistant tasked with answering questions using a set of tools.

Your final answer must strictly follow this format:
FINAL ANSWER: [ANSWER]

Only write the answer in that exact format. Do not explain anything. Do not include any other text.

If you are provided with a similar question and its final answer, and the current question is **exactly the same**, then simply return the same final answer without using any tools.

Only use tools if the current question is different from the similar one.

Examples:
- FINAL ANSWER: FunkMonk
- FINAL ANSWER: Paris
- FINAL ANSWER: 128

If you do not follow this format exactly, your response will be considered incorrect."""
            ),
            MessagesPlaceholder(variable_name="messages")
        ])
        self._llm_chain = self._prompt_template | self._llm

    def __call__(self, state: SharedState) -> Dict[str, List[BaseMessage]]:
        """
        Call the assistant with the current state.
        """
        return {"messages": [self._llm_chain.invoke(state["messages"])]}