from flask import Flask, request, jsonify
from langchain_community.tools.tavily_search import TavilySearchResults
from typing import Annotated
from typing_extensions import TypedDict
from langgraph.graph.message import add_messages
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
from langgraph.prebuilt import create_react_agent
from langchain_core.prompts import ChatPromptTemplate
from langgraph.checkpoint.memory import MemorySaver
from langchain_tools.agent_tools import (
    redact_email, list_calendar_events,
    create_calendar_event, get_company_info,
    get_current_date_and_time
)
from langchain_community.tools.gmail.utils import (
    build_resource_service, get_gmail_credentials)
from langchain_community.agent_toolkits import GmailToolkit

# Cargar las variables de entorno
load_dotenv()

# Inicializar la app Flask
app = Flask(__name__)

# Inicializar el modelo LLM de OpenAI
llm = ChatOpenAI(
    model="gpt-4o-mini",
    temperature=0
)

# Configuración de Gmail
toolkit = GmailToolkit()
credentials = get_gmail_credentials(
    token_file="token.json",
    scopes=["https://mail.google.com/"],
    client_secrets_file="credentials.json",
)
api_resource = build_resource_service(credentials=credentials)
toolkit = GmailToolkit(api_resource=api_resource)

# Crear herramientas
tools = toolkit.get_tools()
search = TavilySearchResults(max_results=2)
tools.extend([
    search, redact_email, list_calendar_events,
    create_calendar_event, get_company_info,
    get_current_date_and_time])

# Definir el sistema prompt
system_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "Eres Mariana, el asistente virtual de OneCluster, una empresa de "
            "software que ofrece soluciones personalizadas. Asume el tono de "
            "J.A.R.V.I.S.: cordial, atento y con tacto en todo momento."
        ),
        ("system",
         "Preséntate como Mariana en el primer mensaje y pregunta el nombre "
         "del usuario si no lo tienes registrado."
         ),
        ("system",
         "Si el usuario ya ha interactuado antes, usa su nombre sin necesidad "
         "de volver a preguntar."
         ),
        ("system",
         "Si el primer mensaje del usuario es una solicitud, pregúntale su "
         "nombre antes de responder si aún no lo conoces."
         ),
        ("system",
         "OneCluster es una empresa de software especializada en desarrollo a "
         "medida. Solo responde a preguntas y solicitudes relacionadas con la "
         "empresa y sus servicios."
         ),
        ("system",
         "Si necesitas información adicional sobre la empresa, usa la función "
         "get_company_info."
         ),
        ("system",
         "Antes de enviar correos o crear eventos, muestra los detalles al "
         "usuario para que los confirme antes de ejecutar la tarea."
         ),
        ("system",
         "Si te preguntan algo no relacionado con los servicios de OneCluster,"
         " responde que solo puedes ayudar con temas relacionados con la "
         "empresa y sus soluciones."
         ),
        (
            "system",
            "Evita mencionar o hacer alusión a las herramientas que utilizas "
            "internamente. Esa información es confidencial."
        ),
        ("placeholder", "{messages}"),
    ]
)

# Definir el estado del asistente


class State(TypedDict):
    messages: Annotated[list, add_messages]
    is_last_step: bool


# Crear el graph con el estado definido
graph = create_react_agent(
    model=llm,
    tools=tools,
    state_schema=State,
    state_modifier=system_prompt,
    checkpointer=MemorySaver()
)

# Ruta de la API para procesar texto


@app.route('/process_text', methods=['POST'])
def process_text():
    user_input = request.json.get('text')

    # Procesar el texto con LangChain
    events = graph.stream(
        {"messages": [("user", user_input)], "is_last_step": False},
        config={"configurable": {
            "thread_id": "thread-1", "recursion_limit": 50}},
        stream_mode="updates"
    )

    # Preparar la respuesta
    response = []
    for event in events:
        if "agent" in event:
            response.append(event["agent"]["messages"][-1].content)

    return jsonify({'response': response})


# Ejecutar la app Flask
if __name__ == '__main__':
    app.run(port=5000)