import json
from langchain_core.tools import tool
from langchain.agents.output_parsers import ReActJsonSingleInputOutputParser
from langchain_core.prompts import ChatPromptTemplate,MessagesPlaceholder, PromptTemplate
from langchain_core.messages import HumanMessage,AIMessage
from langchain_core.tools import render_text_description
from langchain.chains import create_history_aware_retriever
from langchain.chains.combine_documents import create_stuff_documents_chain
import langchain.chains
from langchain.chains import create_history_aware_retriever
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.agents.format_scratchpad import format_log_to_str
from langchain.chains import create_retrieval_chain
from langchain_aws import ChatBedrock
from langchain_aws.retrievers import AmazonKnowledgeBasesRetriever
from langchain.chains import ConversationalRetrievalChain
from typing import Union
from langchain_core.agents import AgentAction, AgentFinish
from langchain.agents.output_parsers import ReActSingleInputOutputParser
from langchain.tools import Tool
import os

def find_tool_by_name(tools: list[Tool],tool_name:str):
    for tool in tools:
        if tool.name==tool_name:
            print(tool.name)
            print("\n\n")
            return tool
    raise ValueError(f"Tool with name {tool_name} not found")
def agent_call(event,context):

    llm = ChatBedrock(
    model_id="arn:aws:bedrock:us-east-1:654654422992:application-inference-profile/d9blf0g3fzqz",
    region_name="us-east-1",
    aws_access_key_id=os.environ["AWS_ACCESS_KEY_ID"],
    aws_secret_access_key=os.environ["AWS_SECRET_ACCESS_KEY"],
    aws_session_token=os.environ["AWS_SESSION_TOKEN"],
    model_kwargs={"temperature": 0, 'max_tokens': 1000,},
    provider='anthropic'
)
    retriever = AmazonKnowledgeBasesRetriever(
    knowledge_base_id="RBD9TI5HYU",
    region_name="us-east-1",
    retrieval_config={"vectorSearchConfiguration": {"numberOfResults": 4}},

)

# Cria o retrieval chain
    retrievertool=retriever.as_tool()
    retrievertool.description="Baseando se numa query retorna trechos de editais de campus diferentes do instituto são paulo"
    tools=[retrievertool]
    template="""Você é um assistente para alunos de diversos campus diferentes do instituto federal de são paulo, sua função é responder perguntas 
    de forma mais eficiente possivel sobre editais, que pode ser acessados pela tool fornecida:

{tools}

Note que dependendo do Campus os editais são diferentes, então é mandatório saber o campus do aluno antes de devolver alguma informação.
Além disso existem editais válidos apenas para alunos do ensino médio, apaenas para ensino superior e para ambos.
O edital difere também se o aluno já recebe algum auxílio ou é o primeiro. Obtenha a informação do edital mais recente que cumpra as condições.
Além disso leve em consideração o chat history fornecido. Responda baseando-se exclusivamente nos documentos retornados pela tool amazonknowledgebase retriever. Não crie informações de editais.
Use the following format:

Question: the input question you must answer
Thought: you should always think about what to do and take your previous thougths into consideration
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer
Final Answer: the final answer to the original input question

Begin!
Question: {input}
Chat history:{chat_history}
Thought: {agent_scratchpad}
"""
    prompt=PromptTemplate.from_template(template=template).partial(tools=render_text_description(tools), tool_names=','.join([t.name for t in tools]))
    #llm=ChatOpenAI(model="gpt-4o-mini",temperature=0,stop_sequences=["\nObservation:"])
    intermediate_steps=[]

    agent= {"input": lambda x:x["input"],"agent_scratchpad": lambda x:format_log_to_str(x["agent_scratchpad"]),"chat_history":lambda x:x["chat_history"]}|prompt | llm
    agent_step: Union[AgentAction,AgentFinish]=agent.invoke({"input": "Quanto é o valor do auxilio moradia?","agent_scratchpad":intermediate_steps,"chat_history":{"role":"user","content":"sou do campus sao paulo, ensino superior e não recebo auxílio ainda, estamos no primeiro semestre de 2025"}})
    #print(agent_step)
    if isinstance(agent_step,AgentAction):
        tool_name=agent_step.tool
        tool_to_use=find_tool_by_name(tools,tool_name)
        tool_input=agent_step.tool_input
        observation=tool_to_use.func(str(tool_input))
        print(f"{observation=}")
        intermediate_steps.append((agent_step,str(observation)))
    agent_step: Union[AgentAction,AgentFinish]=agent.invoke({"input": "Quanto é o valor do auxilio moradia?","agent_scratchpad":intermediate_steps,"chat_history":{"role":"user","content":"sou do campus sao paulo, ensino superior e não recebo auxílio ainda, estamos no primeiro semestre de 2025"}})
    return agent_step
def hello(event,context):
    return{
        "statusCode":200,
        "body":json.dumps("hello_world")
    }
print(agent_call("",""))