import os
import sys
import torch
import uvicorn
import redis
import numpy as np
import random
from fastapi import FastAPI, Query, BackgroundTasks
from fastapi.responses import HTMLResponse
from starlette.middleware.cors import CORSMiddleware
from datasets import load_dataset
from transformers import AutoTokenizer, GPT2LMHeadModel, pipeline
from loguru import logger
from dotenv import load_dotenv
from sklearn.metrics.pairwise import cosine_similarity
from kaggle.api.kaggle_api_extended import KaggleApi

# Importar la librería de spaces
import spaces

sys.path.append('..')

load_dotenv()

huggingface_token = os.getenv('HUGGINGFACE_TOKEN')
kaggle_username = os.getenv('KAGGLE_USERNAME')
kaggle_key = os.getenv('KAGGLE_KEY')

redis_host = os.getenv('REDIS_HOST', 'localhost')
redis_port = os.getenv('REDIS_PORT', 6379)
redis_password = os.getenv('REDIS_PASSWORD', 'huggingface_spaces')
redis_client = redis.Redis(host=redis_host, port=redis_port, password=redis_password, decode_responses=True)

MAX_ITEMS_PER_TABLE = 10000

# Decorador para usar GPU en Spaces
@spaces.GPU()
def generate_responses_gpu(q):
    generated_responses = []
    try:
        for model_name in redis_client.hkeys("models"):
            try:
                model_data = redis_client.hget("models", model_name)
                model = GPT2LMHeadModel.from_pretrained(model_name)
                tokenizer = AutoTokenizer.from_pretrained(model_name)
                
                # Generar valores aleatorios para top_p, top_k y temperature
                top_p = round(random.uniform(0.01, 0.99), 2)
                top_k = random.randint(1, 99)
                temperature = round(random.uniform(0.01, 1.99), 2)

                text_generation_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
                generated_response = text_generation_pipeline(q, do_sample=True, max_length=50, num_return_sequences=5, 
                                                            top_p=top_p, top_k=top_k, temperature=temperature)
                generated_responses.extend([response['generated_text'] for response in generated_response])
            except Exception as e:
                logger.error(f"Error generating response with model {model_name}: {e}")

        if generated_responses:
            similarities = calculate_similarity(q, generated_responses)
            most_coherent_response = generated_responses[np.argmax(similarities)]
            store_to_redis_table(q, "\n".join(generated_responses))
            redis_client.hset("responses", q, most_coherent_response)
        else:
            logger.warning("No valid responses generated.")
    except Exception as e:
        logger.error(f"General error in autocomplete: {e}")


def get_current_table_index():
    return int(redis_client.get("current_table_index") or 0)

def increment_table_index():
    current_index = get_current_table_index()
    redis_client.set("current_table_index", current_index + 1)

def store_to_redis_table(key, content):
    current_index = get_current_table_index()
    table_name = f"table_{current_index}"
    item_count = redis_client.hlen(table_name)
    if item_count >= MAX_ITEMS_PER_TABLE:
        increment_table_index()
        table_name = f"table_{get_current_table_index()}"
    redis_client.hset(table_name, key, content)

def load_and_store_models(model_names):
    for name in model_names:
        try:
            model = GPT2LMHeadModel.from_pretrained(name)
            tokenizer = AutoTokenizer.from_pretrained(name)
            sample_text = "Sample input"
            generated_text = model.generate(tokenizer.encode(sample_text, return_tensors="pt"), max_length=50)
            decoded_text = tokenizer.decode(generated_text[0], skip_special_tokens=True)
            store_to_redis_table(name, decoded_text)
            redis_client.hset("models", name, decoded_text)
        except Exception as e:
            logger.error(f"Error loading model {name}: {e}")

def load_kaggle_datasets(dataset_names):
    api = KaggleApi()
    api.authenticate()
    for dataset_name in dataset_names:
        try:
            api.dataset_download_files(dataset_name, path='./kaggle_datasets', unzip=True)
            dataset = load_dataset('csv', data_files=[f'./kaggle_datasets/{dataset_name}/*.csv'])['train']
            sample_data = dataset.to_pandas().head(10).to_json(orient='records')
            store_to_redis_table(dataset_name, sample_data)
            redis_client.hset("kaggle_datasets", dataset_name, sample_data)
        except Exception as e:
            logger.error(f"Error loading Kaggle dataset {dataset_name}: {e}")


app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"]
)

message_history = []

@app.get('/')
async def index():
    chat_history = redis_client.hgetall(f"table_{get_current_table_index()}")
    chat_history_html = "".join(f"<div class='bot-message'>{msg}</div>" for msg in chat_history.values())

    html_code = f"""
    <!DOCTYPE html>
    <html lang="en">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>ChatGPT Chatbot</title>
        <style>
            body {{ font-family: Arial, sans-serif; margin: 0; padding: 0; background-color: #f4f4f4; }}
            .container {{ max-width: 800px; margin: auto; padding: 20px; }}
            .chat-container {{ background-color: #fff; border-radius: 8px; box-shadow: 0 0 10px rgba(0, 0, 0, 0.1); overflow: hidden; margin-bottom: 20px; }}
            .chat-box {{ height: 300px; overflow-y: auto; padding: 10px; }}
            .chat-input {{ width: calc(100% - 20px); border: none; border-top: 1px solid #ddd; padding: 10px; font-size: 16px; outline: none; }}
            .user-message, .bot-message {{ margin-bottom: 10px; padding: 8px 12px; border-radius: 8px; max-width: 70%; word-wrap: break-word; }}
            .user-message {{ background-color: #007bff; color: #fff; align-self: flex-end; }}
            .bot-message {{ background-color: #4CAF50; color: #fff; }}
            #autocomplete-suggestions {{
                position: absolute;
                background-color: #fff;
                border: 1px solid #ccc;
                border-radius: 4px;
                z-index: 10;
                max-width: calc(100% - 40px);
            }}
            .suggestion {{
                padding: 8px;
                cursor: pointer;
            }}
            .suggestion:hover {{
                background-color: #f0f0f0;
            }}
        </style>
    </head>
    <body>
        <div class="container">
            <h1 style="text-align: center;">ChatGPT Chatbot</h1>
            <div class="chat-container" id="chat-container">
                <div class="chat-box" id="chat-box">
                    {chat_history_html}
                </div>
                <input type="text" class="chat-input" id="user-input" placeholder="Type your message..." autocomplete="off">
                <div id="autocomplete-suggestions"></div>
            </div>
        </div>
        <script>
            const userInput = document.getElementById('user-input');
            const autocompleteSuggestions = document.getElementById('autocomplete-suggestions');

            userInput.addEventListener('keyup', function(event) {{
                if (event.key === 'Enter') {{
                    event.preventDefault();
                    sendMessage();
                }} else {{
                    fetch(`/autocomplete?q=` + encodeURIComponent(userInput.value))
                    .then(response => response.json())
                    .then(data => {{
                        displayAutocompleteSuggestions(data.suggestions);
                    }})
                    .catch(error => {{
                        console.error('Error:', error);
                    }});
                }}
            }});

            function displayAutocompleteSuggestions(suggestions) {{
                autocompleteSuggestions.innerHTML = '';
                if (suggestions.length > 0) {{
                    suggestions.forEach(suggestion => {{
                        const suggestionElement = document.createElement('div');
                        suggestionElement.className = 'suggestion';
                        suggestionElement.innerText = suggestion;
                        suggestionElement.onclick = () => {{
                            userInput.value = suggestion;
                            autocompleteSuggestions.innerHTML = '';
                        }};
                        autocompleteSuggestions.appendChild(suggestionElement);
                    }});
                }}
            }}

            function sendMessage() {{
                const userMessage = userInput.value.trim();
                if (userMessage === '') return;

                appendMessage('user', userMessage);
                userInput.value = '';
                autocompleteSuggestions.innerHTML = '';

                fetch(`/autocomplete?q=` + encodeURIComponent(userMessage))
                .then(response => response.json())
                .then(data => {{
                    fetch(`/get_response?q=` + encodeURIComponent(userMessage))
                    .then(response => response.json())
                    .then(data => {{
                        const botMessage = data.response;
                        appendMessage('bot', botMessage);
                    }})
                    .catch(error => {{
                        console.error('Error:', error);
                    }});
                }})
                .catch(error => {{
                    console.error('Error:', error);
                }});
            }}

            function appendMessage(sender, message) {{
                const chatBox = document.getElementById('chat-box');
                const messageElement = document.createElement('div');
                messageElement.className = sender + '-message';
                messageElement.innerText = message;
                chatBox.appendChild(messageElement);
            }}
        </script>
    </body>
    </html>
    """
    return HTMLResponse(content=html_code, status_code=200)

def calculate_similarity(base_text, candidate_texts):
    base_vector = np.array([len(base_text)])
    similarities = []
    for text in candidate_texts:
        candidate_vector = np.array([len(text)])
        similarity = cosine_similarity([base_vector], [candidate_vector])
        similarities.append(similarity[0][0])
    return similarities

@app.get('/autocomplete')
async def autocomplete(q: str = Query(..., title='query'), background_tasks: BackgroundTasks = BackgroundTasks()):
    global message_history
    message_history.append(('user', q))

    suggestions = []
    if q:
        for key in redis_client.hkeys("responses"):
            if q.lower() in key.lower():
                suggestions.append(key)
    
    # Lanzar la tarea en segundo plano utilizando la función decorada con @spaces.GPU()
    background_tasks.add_task(generate_responses_gpu, q) 
    return {"status": "Processing request, please wait...", "suggestions": suggestions}

@app.get('/get_response')
async def get_response(q: str = Query(..., title='query')):
    response = redis_client.hget("responses", q)
    return {"response": response}



if __name__ == '__main__':
    gpt2_models = [
        "gpt2",
        "gpt2-medium",
        "gpt2-large",
        "gpt2-xl"
    ]

    programming_models = [
        "google/bert2bert_L-24_uncased",
        "microsoft/CodeGPT-small-java",
        "microsoft/CodeGPT-small-python",
        "Salesforce/codegen-350M-multi"
    ]

    kaggle_datasets = [
        "uciml/iris",
        "arshid/iris-flower-dataset",
        "heesoo37/120-years-of-olympic-history-athletes-and-results"
    ]

    load_and_store_models(gpt2_models + programming_models)
    load_kaggle_datasets(kaggle_datasets)

    uvicorn.run(app=app, host='0.0.0.0', port=int(os.getenv("PORT", 7860)))