import os
import subprocess
import threading
import time
import asyncio
import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from langchain.prompts import PromptTemplate
from langchain_community.chat_models import ChatOllama
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
from langchain_core.output_parsers import JsonOutputParser, StrOutputParser
from langgraph.graph import END, StateGraph
from typing_extensions import TypedDict
from fastapi.responses import StreamingResponse
from ollama import AsyncClient
import gc
import psutil
import torch
from functools import lru_cache

# Asegúrate de tener la librería Ollama instalada
# pip install ollama langchain langchain_community langgraph

# Ruta de Ollama
OLLAMA = os.path.expanduser("~/ollama")
if not os.path.exists(OLLAMA):
    print("Ollama no encontrado, descargando...")
    subprocess.run("curl -L https://ollama.com/download/ollama-linux-amd64 -o ~/ollama", shell=True)
    os.chmod(OLLAMA, 0o755)

# Iniciar el servidor Ollama en segundo plano
async def ollama_service_thread():
    print("Iniciando el servicio de Ollama")
    subprocess.run("~/ollama serve -1", shell=True)  # -1 indica mantener el servidor en ejecución indefinidamente

# Iniciar el hilo de servicio Ollama
print("Creando y comenzando el hilo del servicio Ollama")
OLLAMA_SERVICE_THREAD = threading.Thread(target=asyncio.run, args=(ollama_service_thread(),))
OLLAMA_SERVICE_THREAD.start()

# Esperar a que Ollama se inicie
print("Esperando a que Ollama inicie...")
time.sleep(10)

# Descargar el modelo de Hugging Face si no está disponible
async def download_ollama_model(model_name='hf.co/MaziyarPanahi/Llama-3.2-3B-Instruct-uncensored-GGUF:IQ1_S'):
    try:
        print(f"Descargando el modelo: {model_name}")
        subprocess.run(["ollama", "pull", model_name], check=True)
    except subprocess.CalledProcessError as e:
        print(f"Error al descargar el modelo: {e}")
        raise

# Descargar el modelo de Ollama en el hilo principal
download_ollama_model("hf.co/MaziyarPanahi/Llama-3.2-3B-Instruct-uncensored-GGUF:IQ1_S")

# Crear la aplicación FastAPI
app = FastAPI()

# Definir el modelo de datos para recibir las consultas en la API
class QueryRequest(BaseModel):
    query: str

# Definir el modelo de lenguaje de Ollama
local_llm = 'hf.co/MaziyarPanahi/Llama-3.2-3B-Instruct-uncensored-GGUF:IQ1_S'
llama3 = ChatOllama(model=local_llm)

# Definir la herramienta de búsqueda web usando DuckDuckGo
wrapper = DuckDuckGoSearchAPIWrapper(max_results=1)
web_search_tool = DuckDuckGoSearchRun(api_wrapper=wrapper)

# Implementar optimización extrema mediante LRU Cache
@lru_cache(maxsize=1024)  # Cache de tamaño limitado para resultados previos
async def cached_search(query):
    return await web_search_tool.invoke(query)

# Definir los prompts para generación y enrutamiento
generate_prompt = PromptTemplate(
    template="""
    <|begin_of_text|>
    <|start_header_id|>system<|end_header_id|> 
    You are an AI assistant for Research Question Tasks, that synthesizes web search results. 
    Strictly use the following pieces of web search context to answer the question. If you don't know the answer, just say that you don't know. 
    Keep the answer concise, but provide all of the details you can in the form of a research report. 
    Only make direct references to material if provided in the context.
    <|eot_id|>
    <|start_header_id|>user<|end_header_id|> 
    Question: {question} 
    Web Search Context: {context} 
    Answer: 
    <|eot_id|> 
    <|start_header_id|>assistant<|end_header_id|>""",
    input_variables=["question", "context"],
)

generate_chain = generate_prompt | llama3 | StrOutputParser()

router_prompt = PromptTemplate(
    template="""
    <|begin_of_text|>
    <|start_header_id|>system<|end_header_id|> 
    You are an expert at routing a user question to either the generation stage or web search. 
    Use the web search for questions that require more context for a better answer, or recent events.
    Otherwise, you can skip and go straight to the generation phase to respond.
    You do not need to be stringent with the keywords in the question related to these topics.
    Give a binary choice 'web_search' or 'generate' based on the question. 
    Return the JSON with a single key 'choice' with no premable or explanation. 
    Question to route: {question} 
    <|eot_id|> 
    <|start_header_id|>assistant<|end_header_id|>""",
    input_variables=["question"],
)

question_router = router_prompt | llama3 | JsonOutputParser()

query_prompt = PromptTemplate(
    template="""
    <|begin_of_text|>
    <|start_header_id|>system<|end_header_id|> 
    You are an expert at crafting web search queries for research questions.
    More often than not, a user will ask a basic question that they wish to learn more about, however it might not be in the best format. 
    Reword their query to be the most effective web search string possible.
    Return the JSON with a single key 'query' with no premable or explanation. 
    Question to transform: {question} 
    <|eot_id|> 
    <|start_header_id|>assistant<|end_header_id|>""",
    input_variables=["question"],
)

query_chain = query_prompt | llama3 | JsonOutputParser()

# Definir el estado del grafo
class GraphState(TypedDict):
    question: str
    generation: str
    search_query: str
    context: str

# Nodos de procesamiento
async def generate(state):
    print("Step: Generating Final Response")
    question = state["question"]
    context = state["context"]
    generation = await generate_chain.invoke({"context": context, "question": question})
    return {"generation": generation}

async def transform_query(state):
    print("Step: Optimizing Query for Web Search")
    question = state['question']
    gen_query = await query_chain.invoke({"question": question})
    search_query = gen_query.get("query", "")  # Asegurarnos de que estamos obteniendo la clave correcta
    return {"search_query": search_query}

async def web_search(state):
    search_query = state['search_query']
    print(f'Step: Searching the Web for: "{search_query}"')
    try:
        # Se usa la caché para optimizar los resultados
        search_result = await cached_search(search_query)
        if isinstance(search_result, str):  # Si la respuesta es una cadena, la convertimos en un diccionario
            print(f"Respuesta de búsqueda web es cadena: {search_result}")
            return {"context": search_result}
        elif isinstance(search_result, dict):  # Si es un diccionario, lo usamos directamente
            return {"context": search_result}
        else:
            raise ValueError("Respuesta de búsqueda web no es válida")
    except Exception as e:
        print(f"Web search failed: {e}")
        return None  # Si la búsqueda falla, no devuelve contexto

async def route_question(state):
    print("Step: Routing Query")
    question = state['question']
    output = await question_router.invoke({"question": question})
    if output.get('choice') == "web_search":
        print("Step: Routing Query to Web Search")
        return "websearch"
    elif output.get('choice') == 'generate':
        print("Step: Routing Query to Generation")
        return "generate"

# Crear el grafo de estado
workflow = StateGraph(GraphState)
workflow.add_node("websearch", web_search)
workflow.add_node("transform_query", transform_query)
workflow.add_node("generate", generate)

workflow.set_conditional_entry_point(
    route_question,
    {
        "websearch": "transform_query",
        "generate": "generate",
    },
)

# Gestión de recursos de CPU, RAM y GPU
def release_resources():
    try:
        torch.cuda.empty_cache()
        gc.collect()
    except Exception as e:
        print(f"Failed to release resources: {e}")

def resource_manager():
    MAX_RAM_PERCENT = 1  # Ajustar según sea necesario
    MAX_CPU_PERCENT = 1
    MAX_GPU_PERCENT = 1
    MAX_RAM_MB = 1  # Ajustar según la memoria disponible

    while True:
        try:
            virtual_mem = psutil.virtual_memory()
            current_ram_percent = virtual_mem.percent
            current_ram_mb = virtual_mem.used / (1 * 1)  # Convertir a MB

            if current_ram_percent > MAX_RAM_PERCENT or current_ram_mb > MAX_RAM_MB:
                release_resources()

            current_cpu_percent = psutil.cpu_percent()
            if current_cpu_percent > MAX_CPU_PERCENT:
                psutil.Process(os.getpid()).nice()

            if torch.cuda.is_available():
                gpu = torch.cuda.current_device()
                gpu_mem = torch.cuda.memory_allocated(gpu) / (1024 * 1024)  # Convertir a MB

                if gpu_mem > MAX_GPU_PERCENT:
                    release_resources()

        except Exception as e:
            print(f"Error en el gestor de recursos: {e}")

resource_manager()

# Procesar la consulta en paralelo
async def process_query_in_parallel(query):
    try:
        state = GraphState(question=query, generation="", search_query="", context="")
        return await workflow.invoke(state)
    except Exception as e:
        print(f"Error en la ejecución paralela: {e}")
        raise

# Ruta de la API para manejar consultas
@app.post("/query")
async def query_handler(request: QueryRequest):
    try:
        query = request.query
        result = await process_query_in_parallel(query)  # Llamada asíncrona
        return {"results": result}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Ejecutar el servidor FastAPI
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)