from fastapi import FastAPI, HTTPException, Request
import uvicorn
import requests
import os
import io
import asyncio
from typing import List, Dict, Any
from tqdm import tqdm
from llama_cpp import Llama
import aiofiles
import time
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

app = FastAPI()

model_configs = [
    {"repo_id": "Ffftdtd5dtft/gpt2-xl-Q2_K-GGUF", "filename": "gpt2-xl-q2_k.gguf", "name": "GPT-2 XL"},
    {"repo_id": "Ffftdtd5dtft/Meta-Llama-3.1-8B-Instruct-Q2_K-GGUF", "filename": "meta-llama-3.1-8b-instruct-q2_k.gguf", "name": "Meta Llama 3.1-8B Instruct"},
    {"repo_id": "Ffftdtd5dtft/gemma-2-9b-it-Q2_K-GGUF", "filename": "gemma-2-9b-it-q2_k.gguf", "name": "Gemma 2-9B IT"},
    {"repo_id": "Ffftdtd5dtft/gemma-2-27b-Q2_K-GGUF", "filename": "gemma-2-27b-q2_k.gguf", "name": "Gemma 2-27B"},
    {"repo_id": "Ffftdtd5dtft/Phi-3-mini-128k-instruct-Q2_K-GGUF", "filename": "phi-3-mini-128k-instruct-q2_k.gguf", "name": "Phi-3 Mini 128K Instruct"},
    {"repo_id": "Ffftdtd5dtft/Meta-Llama-3.1-8B-Q2_K-GGUF", "filename": "meta-llama-3.1-8b-q2_k.gguf", "name": "Meta Llama 3.1-8B"},
    {"repo_id": "Ffftdtd5dtft/Qwen2-7B-Instruct-Q2_K-GGUF", "filename": "qwen2-7b-instruct-q2_k.gguf", "name": "Qwen2 7B Instruct"},
    {"repo_id": "Ffftdtd5dtft/starcoder2-3b-Q2_K-GGUF", "filename": "starcoder2-3b-q2_k.gguf", "name": "Starcoder2 3B"},
    {"repo_id": "Ffftdtd5dtft/Qwen2-1.5B-Instruct-Q2_K-GGUF", "filename": "qwen2-1.5b-instruct-q2_k.gguf", "name": "Qwen2 1.5B Instruct"},
    {"repo_id": "Ffftdtd5dtft/starcoder2-15b-Q2_K-GGUF", "filename": "starcoder2-15b-q2_k.gguf", "name": "Starcoder2 15B"},
    {"repo_id": "Ffftdtd5dtft/gemma-2-2b-it-Q2_K-GGUF", "filename": "gemma-2-2b-it-q2_k.gguf", "name": "Gemma 2-2B IT"},
    {"repo_id": "Ffftdtd5dtft/sarvam-2b-v0.5-Q2_K-GGUF", "filename": "sarvam-2b-v0.5-q2_k.gguf", "name": "Sarvam 2B v0.5"},
    {"repo_id": "Ffftdtd5dtft/WizardLM-13B-Uncensored-Q2_K-GGUF", "filename": "wizardlm-13b-uncensored-q2_k.gguf", "name": "WizardLM 13B Uncensored"},
    {"repo_id": "Ffftdtd5dtft/WizardLM-7B-Uncensored-Q2_K-GGUF", "filename": "wizardlm-7b-uncensored-q2_k.gguf", "name": "WizardLM 7B Uncensored"},
    {"repo_id": "Ffftdtd5dtft/Qwen2-Math-7B-Instruct-Q2_K-GGUF", "filename": "qwen2-math-7b-instruct-q2_k.gguf", "name": "Qwen2 Math 7B Instruct"}
]

models_dir = "modelos"
models = {}

class ModelManager:
    def __init__(self):
        self.model_parts = {}
        self.load_lock = asyncio.Lock()
        self.index_lock = asyncio.Lock()
        self.part_size = 1024 * 1024 

    async def download_model(self, model_config):
        model_path = os.path.join(models_dir, model_config['filename'])
        if not os.path.exists(model_path):
            url = f"https://huggingface.co/{model_config['repo_id']}/resolve/main/{model_config['filename']}"
            print(f"Descargando modelo desde {url}")
            try:
                start_time = time.time()
                response = requests.get(url, stream=True)
                response.raise_for_status()

                total_size = int(response.headers.get('content-length', 0))
                with open(model_path, 'wb') as f:
                    with tqdm(total=total_size, unit='B', unit_scale=True, desc=f"Descargando {model_config['filename']}") as pbar:
                        for chunk in response.iter_content(chunk_size=8192):
                            f.write(chunk)
                            pbar.update(len(chunk))
                end_time = time.time()
                download_duration = end_time - start_time
                print(f"Descarga completa para {model_config['name']} en {download_duration:.2f} segundos")
            except requests.RequestException as e:
                raise HTTPException(status_code=500, detail=f"Error al descargar el modelo: {e}")
        else:
            print(f"Modelo {model_config['filename']} ya descargado.")
        return model_path

    async def download_all_models(self):
        async with self.load_lock:
            download_tasks = [self.download_model(config) for config in model_configs]
            await asyncio.gather(*download_tasks)

    async def load_all_models(self):
        async with self.load_lock:
            load_tasks = [self.load_model(config) for config in model_configs]
            await asyncio.gather(*load_tasks)

    async def load_model(self, model_config):
        model_name = model_config['name']
        if model_name not in models:
            try:
                model_path = os.path.join(models_dir, model_config['filename'])
                start_time = time.time()
                print(f"Cargando modelo desde {model_path}")

                llama = Llama(model_path=model_path)

                end_time = time.time()
                load_duration = end_time - start_time
                if load_duration > 0:
                    print(f"Modelo {model_name} tardó {load_duration:.2f} segundos en cargar")
                else:
                    print(f"Modelo {model_name} cargado correctamente en {load_duration:.2f} segundos")

                tokenizer = llama.tokenizer
                models[model_name] = {
                    'model': llama,
                    'tokenizer': tokenizer,
                }
            except Exception as e:
                print(f"Error al cargar el modelo: {e}")

    async def generate_response(self, user_input, model_name=None, top_k=50, top_p=0.95, temperature=0.8):
        results = []
        if model_name:
            model_data = models.get(model_name)
            if not model_data:
                return {"model_name": model_name, "error": "Modelo no encontrado"}
            try:
                tokenizer = model_data['tokenizer']
                input_ids = tokenizer(user_input).input_ids
                outputs = model_data['model'].generate(
                    [input_ids],
                    top_k=top_k,
                    top_p=top_p,
                    temperature=temperature
                )
                generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
                parts = []
                while len(generated_text) > 1000:
                    part = generated_text[:1000]
                    parts.append(part)
                    generated_text = generated_text[1000:]
                parts.append(generated_text)
                results.append({
                    'model_name': model_name,
                    'generated_text': generated_text,
                    'generated_text_parts': parts
                })
            except Exception as e:
                return {'model_name': model_name, 'error': str(e)}
        else:
            for model_name, model_data in models.items():
                try:
                    tokenizer = model_data['tokenizer']
                    input_ids = tokenizer(user_input).input_ids
                    outputs = model_data['model'].generate(
                        [input_ids],
                        top_k=top_k,
                        top_p=top_p,
                        temperature=temperature
                    )
                    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
                    parts = []
                    while len(generated_text) > 1000:
                        part = generated_text[:1000]
                        parts.append(part)
                        generated_text = generated_text[1000:]
                    parts.append(generated_text)
                    results.append({
                        'model_name': model_name,
                        'generated_text': generated_text,
                        'generated_text_parts': parts
                    })
                except Exception as e:
                    results.append({'model_name': model_name, 'error': str(e)})

        if len(results) > 1:
            best_response = self.choose_best_response(user_input, results)
        elif len(results) == 1:
            best_response = results[0]
        else: 
            return {"model_name": "Error", "error": "No se pudo generar una respuesta con ningún modelo."}

        return best_response

    def choose_best_response(self, user_input, responses):
        valid_responses = [r for r in responses if 'error' not in r]

        tfidf = TfidfVectorizer()
        response_texts = [r['generated_text'] for r in valid_responses]
        tfidf_matrix = tfidf.fit_transform([user_input] + response_texts)
        similarities = cosine_similarity(tfidf_matrix[0:1], tfidf_matrix[1:])

        best_index = similarities.argmax()
        best_response = valid_responses[best_index]

        return best_response

@app.post("/generate/")
async def generate(request: Request):
    data = await request.json()
    user_input = data.get('input', '')
    model_name = data.get('model')
    top_k = data.get('top_k', 50)
    top_p = data.get('top_p', 0.95)
    temperature = data.get('temperature', 0.8)
    if not user_input:
        raise HTTPException(status_code=400, detail="Se requiere una entrada de usuario.")

    try:
        response = await model_manager.generate_response(user_input, model_name, top_k, top_p, temperature)
        return {"response": response}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/models")
async def get_available_models():
    return {"models": [config['name'] for config in model_configs]}

async def load_models_on_startup():
    global model_manager
    model_manager = ModelManager()
    await model_manager.download_all_models()
    await model_manager.load_all_models()

@app.on_event("startup")
async def startup_event():
    await load_models_on_startup()
    print("Modelos cargados. API lista.")

if __name__ == "__main__":
    if not os.path.exists(models_dir):
        os.makedirs(models_dir)

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

html_code = """
<!DOCTYPE html>
<html>
<head>
<title>Chatbot</title>
<style>
body {
    display: flex;
    justify-content: center;
    align-items: center;
    height: 100vh;
    margin: 0;
    font-family: sans-serif;
}

.container {
    border: 1px solid #ccc;
    border-radius: 5px;
    width: 400px;
    height: 500px;
    overflow: hidden;
}

.chat-log {
    padding: 10px;
    height: 400px;
    overflow-y: scroll;
}

.chat-message {
    margin-bottom: 10px;
    padding: 8px;
    border-radius: 5px;
}

.user-message {
    background-color: #eee;
}

.bot-message {
    background-color: #ccf;
}

.input-area {
    display: flex;
    padding: 10px;
}

#user-input {
    flex: 1;
    padding: 8px;
    border: 1px solid #ccc;
    border-radius: 5px;
}

#send-button {
    padding: 8px 15px;
    background-color: #4CAF50;
    color: white;
    border: none;
    border-radius: 5px;
    cursor: pointer;
    margin-left: 10px;
}

#model-select {
    width: 100%;
    padding: 8px;
    border: 1px solid #ccc;
    border-radius: 5px;
    margin-bottom: 10px;
}
</style>
</head>
<body>
    <div class="container">
        <div class="chat-log" id="chat-log">
        </div>
        <div class="input-area">
            <input type="text" id="user-input" placeholder="Escribe tu mensaje...">
            <button id="send-button">Enviar</button>
        </div>
        <select id="model-select">
            <option value="">Todos los modelos</option>
            </select> 
    </div>

    <script>
        const chatLog = document.getElementById('chat-log');
        const userInput = document.getElementById('user-input');
        const sendButton = document.getElementById('send-button');
        const modelSelect = document.getElementById('model-select');
        let currentConversationId = null;

        async function startNewConversation() {
        }

        startNewConversation();

        async function getAvailableModels() {
            const response = await fetch('/models');
            const data = await response.json();
            return data.models;
        }

        async function displayAvailableModels() {
            const models = await getAvailableModels();
            models.forEach(model => {
                const option = document.createElement('option');
                option.value = model;
                option.text = model;
                modelSelect.add(option);
            });
        }

        displayAvailableModels();

        sendButton.addEventListener('click', async () => {
            const userMessage = userInput.value;
            userInput.value = '';
            const selectedModel = modelSelect.value;

            appendMessage('user', userMessage);

            const response = await fetch('/generate/', {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json'
                },
                body: JSON.stringify({ input: userMessage, model: selectedModel }) 
            });

            const data = await response.json();
            if (data.response.error) {
                appendMessage('bot', `Error del modelo ${data.response.model_name}: ${data.response.error}`);
            } else {
                data.response.generated_text_parts.forEach(part => {
                    appendMessage('bot', part);
                });
            }
        });

        function appendMessage(role, message) {
            const messageElement = document.createElement('div');
            messageElement.classList.add('chat-message', `${role}-message`);
            messageElement.textContent = message;
            chatLog.appendChild(messageElement);
            chatLog.scrollTop = chatLog.scrollHeight;
        }
    </script>
</body>
</html>
"""