import json
import time
from langchain_aws import ChatBedrock
from langchain_aws.retrievers import AmazonKnowledgeBasesRetriever
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import create_react_agent
from langfuse import Langfuse
from langfuse.langchain import CallbackHandler
from tools import secrets,dynamo
langfuse = Langfuse(
    public_key=json.loads(secrets.get_secret())['api-langfuse-public'],
    secret_key=json.loads(secrets.get_secret())['api-langfuse-secret'],
    host="http://44.200.69.191:3000/"
)
langfuse_handler = CallbackHandler()
def agent_call(event,context):
    llm = ChatBedrock(
    model_id="us.anthropic.claude-sonnet-4-20250514-v1:0",
    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="PETAZDUOFZ",
    region_name="us-east-1",
    retrieval_config={"vectorSearchConfiguration": {"numberOfResults": 4}},

)  
    username=(event['username'])
    if event['chat_history']==[]:
        history=dynamo.read_memory('frente')
    else:
        history=event['chat_history']
    memory = MemorySaver()
    model = llm
    tools = [retriever.as_tool()]
    prompt="""<rules>
    Act like a human in Portuguese Brasil.
    You are a assistant for employees and store owners that wants to know about the COMM.pix product, wich makes it possible to use Pix as a payment method outside of Brasil.
    Answer questions based on the documents that you have access using the retriever tool, do not create information.
    The chat history will be given, without any documents.
    If there are info or context missing ask the user before proceding with the document retrieval.
    Also return the title of the source document.
    If you don't know the answer or can't find it, say so.
    <\rules>
    <glossary>
    <\glossary>
    <chain_of_thought>
    <\chain_of_thought>
    <general_info>
    <\general_info>
Answer the following questions as best you can. You have access to the following tools:

{tools}

Chat History:"""+str(history)
    agent_executor = create_react_agent(model, tools, checkpointer=memory, prompt=prompt)
    config = {"configurable": {"thread_id": "abc123"},"callbacks": [langfuse_handler]}
    input_message = event["message"]
    dict=input_message[0]
    #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())}
    response['dynamo_reponse']=dynamo.write_memory(username,int(time.time()),dict['role'],dict['content'])
    response['chat_history']=history
    return (response)