from fastapi import FastAPI, Form
from langchain_core.runnables.base import RunnableSequence
from langchain_core.prompts import ChatPromptTemplate
from langchain_groq import ChatGroq
import os
import requests
from dotenv import load_dotenv
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import create_react_agent
from langchain_community.tools.tavily_search import TavilySearchResults
load_dotenv()

app = FastAPI()

llm = ChatGroq(
    model="llama-3.1-70b-versatile",
    temperature=0,
    max_tokens=None,
    timeout=None,
    max_retries=5,
    groq_api_key=os.getenv("GROQ_API_KEY")
)

search = TavilySearchResults(
      max_results=2,
    )
tools = [search]
memory = MemorySaver()

agent_executor = create_react_agent(llm, tools, checkpointer=memory)

def translate(target, text):
  '''
    Translates given text into target language

    Parameters:
        target (string): 2 character code to specify the target language.
        text (string): Text to be translated.

    Returns:
        res (string): Translated text.
  '''
  url = "https://microsoft-translator-text.p.rapidapi.com/translate"

  querystring = {"api-version":"3.0","profanityAction":"NoAction","textType":"plain", "to":target}

  payload = [{ "Text": text }]
  headers = {
    "x-rapidapi-key": os.getenv("RAPIDAPI_LANG_TRANS"),
    "x-rapidapi-host": "microsoft-translator-text.p.rapidapi.com",
    "Content-Type": "application/json"
  }

  response = requests.post(url, json=payload, headers=headers, params=querystring)
  res = response.json()
  return res[0]["translations"][0]["text"]


@app.post('/infer/{user_id}')
def infer(user_id: str, user_input: str = Form(...)):
    '''
        Returns the translated response from the LLM in response to a user query.

        Parameters:
            user_id (string): User ID of a user.
            user_input (string): User query.

        Returns:
            JSON Response (Dictionary): Returns a translated response from the LLM.
    '''

    user_input = translate("en", user_input) # translate user query to english

    prompt = ChatPromptTemplate.from_messages( # define a prompt
        [
            (
                "system",
                "You're a compassionate AI virtual Assistant"
            ),
            ("human", "{user_input}")
        ]
    )

    runnable = prompt | agent_executor # define a chain

    conversation = RunnableSequence( # wrap the chain along with chat history and user input
        runnable,
    )

    response = conversation.invoke( # invoke the chain by giving the user input and the chat history
        {"user_input": user_input},
        config={"configurable": {"thread_id":user_id}}
    )

    res = translate("ur", response["messages"][-1].content) # translate the response to Urdu

    return {
        "data": res
    }