from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, RedirectResponse
from langchain_openai import ChatOpenAI
from langserve import add_routes

from langchain_core.prompts import ChatPromptTemplate

from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.tools.gmail.utils import build_resource_service
from langchain_community.agent_toolkits import GmailToolkit

from app.langchain_tools.agent_tools import (
    redact_email,
    get_company_info,
    get_current_date_and_time,
)

from langgraph.graph.message import add_messages
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import MemorySaver

from typing import Annotated
from typing_extensions import TypedDict
from dotenv import load_dotenv
import os

load_dotenv()

app = FastAPI()
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.9)


# # Crear herramientas
tools = []
search = TavilySearchResults(max_results=2)
tools.extend(
    [
        search,
        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(),
)


@app.get("/")
async def redirect_root_to_docs():
    return RedirectResponse("/docs")


# # Edit this to add the chain you want to add
add_routes(app, llm, path="/openai")


@app.post("/process_text")
async def process_text(request: Request):
    data = await request.json()
    user_input = data.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 JSONResponse(content={"response": response})


if __name__ == "__main__":
    config = {"configurable": {"thread_id": "thread-1", "recursion_limit": 50}}

    # Modo interactivo por defecto
    import sys

    if "--server" not in sys.argv:
        while True:
            user_input = input("User: ")
            if user_input.lower() in ["quit", "exit", "q"]:
                print("Goodbye!")
                break

            events = graph.stream(
                {"messages": [("user", user_input)], "is_last_step": False},
                config,
                stream_mode="updates",
            )

            for event in events:
                if "agent" in event:
                    print(f"\nAsistente: {event['agent']['messages'][-1].content}\n")
    else:
        # Modo servidor con uvicorn
        import uvicorn

        uvicorn.run(app, host="0.0.0.0", port=8080)