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
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import create_react_agent
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 hello(event,context):
    llm = ChatBedrock(
    model_id="arn:aws:bedrock:us-east-1:654654422992:application-inference-profile/glvyppv11ds5",
    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.1, 'max_tokens': 1000,},
    provider='anthropic'
)
    retriever = AmazonKnowledgeBasesRetriever(
    knowledge_base_id="RBD9TI5HYU",
    region_name="us-east-1",
    retrieval_config={"vectorSearchConfiguration": {"numberOfResults": 4}},

)
    memory = MemorySaver()
    model = llm
    tools = [retriever.as_tool()]
    prompt="""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. 
        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.  Não crie informações de editais, apenas faça perguntas para o contexto e busque dos editais quando estiver com informações suficientes. Retorne o titulo do documento consultado, não o nome do arquivo. 
Documentos com a sigla MTO são do campus matão. Documentos com a sigla SPO são do campus são paulo. É imperativo não obeter documentos de outros campus que não sejam o solicitado, não use informações deles. Os valores, se estiverem em formato de tabela, se encontrarão após o tipo da modaldie (moradia,transporte e etc), nunca antes. Não é nescessário buscar informações dos editais se uma dúvida não for apresentada ou se houver uma informação importante faltando.
Use sempre os documentos mais recentes, atualmente estamos no segundo semestre de 2025.
Caso seja mencionado que um dos pais não é registrado isso significa que é trabalhador informal."""
    agent_executor = create_react_agent(model, tools, checkpointer=memory, prompt=prompt)
    config = {"configurable": {"thread_id": "abc123"}}
    input_message = event["message"]
    #input_message=[{"role":"user","content":"aluno superior, nunca recebi auxilio, campus são paulo, Meu pai não é registrado, como faço para ganhar auxilio?"}]
    response=""
    for step in agent_executor.stream({"messages": input_message}, config, stream_mode="values"):
        response={"json":(step["messages"][-1].text())}
    return (response)