from dotenv import load_dotenv
import os
import json
import requests
import redis
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    AutoModelForCausalLM,
)
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from torch.optim import AdamW
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import HTMLResponse
import multiprocessing
import time
import uuid
import random

load_dotenv()

REDIS_HOST = os.getenv('REDIS_HOST')
REDIS_PORT = os.getenv('REDIS_PORT')
REDIS_PASSWORD = os.getenv('REDIS_PASSWORD')

app = FastAPI()

default_language = "es"

class ChatbotService:
    def __init__(self):
        self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)
        self.model_name = "response_model"
        self.tokenizer_name = "response_tokenizer"

    def get_response(self, user_id, message, language=default_language):
        model = self.load_model_from_redis()
        tokenizer = self.load_tokenizer_from_redis()

        if model is None or tokenizer is None:
            return "El modelo aún no está listo. Por favor, inténtelo de nuevo más tarde."

        input_text = f"Usuario: {message} Asistente:"
        input_ids = tokenizer.encode(input_text, return_tensors="pt").to("cpu")

        with torch.no_grad():
            output = model.generate(input_ids=input_ids, max_length=100, num_beams=5, no_repeat_ngram_size=2, early_stopping=True)

        response = tokenizer.decode(output[0], skip_special_tokens=True)
        response = response.replace(input_text, "").strip()

        return response

    def load_model_from_redis(self):
        model_data_bytes = self.redis_client.get(f"model:{self.model_name}")
        if model_data_bytes:
            model = AutoModelForCausalLM.from_pretrained("gpt2")
            model.load_state_dict(torch.load(model_data_bytes))
            return model
        else:
            return None

    def load_tokenizer_from_redis(self):
        tokenizer_data_bytes = self.redis_client.get(f"tokenizer:{self.tokenizer_name}")
        if tokenizer_data_bytes:
            tokenizer = AutoTokenizer.from_pretrained("gpt2")
            tokenizer.add_tokens(json.loads(tokenizer_data_bytes))
            return tokenizer
        else:
            return None

chatbot_service = ChatbotService()

class UnifiedModel(nn.Module):
    def __init__(self, models):
        super(UnifiedModel, self).__init__()
        self.models = nn.ModuleList(models)
        hidden_size = self.models[0].config.hidden_size
        self.projection = nn.Linear(len(models) * 3, 768)
        self.classifier = nn.Linear(hidden_size, 3)

    def forward(self, input_ids, attention_mask):
        hidden_states = []
        for model, input_id, attn_mask in zip(self.models, input_ids, attention_mask):
            outputs = model(
                input_ids=input_id,
                attention_mask=attn_mask
            )
            hidden_states.append(outputs.logits)

        concatenated_hidden_states = torch.cat(hidden_states, dim=1)
        projected_features = self.projection(concatenated_hidden_states)
        logits = self.classifier(projected_features)
        return logits

    @staticmethod
    def load_model_from_redis(redis_client):
        model_name = "unified_model"
        model_data_bytes = redis_client.get(f"model:{model_name}")
        if model_data_bytes:
            model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3)
            model.load_state_dict(torch.load(model_data_bytes))
        else:
            model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3)
        return UnifiedModel([model, model])

class SyntheticDataset(Dataset):
    def __init__(self, tokenizers, data):
        self.tokenizers = tokenizers
        self.data = data

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        item = self.data[idx]
        text = item['text']
        label = item['label']
        tokenized = {}
        for name, tokenizer in self.tokenizers.items():
            tokens = tokenizer(text, padding="max_length", truncation=True, max_length=128)
            tokenized[f"input_ids_{name}"] = torch.tensor(tokens["input_ids"])
            tokenized[f"attention_mask_{name}"] = torch.tensor(tokens["attention_mask"])
        tokenized["labels"] = torch.tensor(label)
        return tokenized

conversation_history = {}

@app.post("/process")
async def process(request: Request):
    data = await request.json()
    redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)

    tokenizers = {}
    models = {}

    model_name = "unified_model"
    tokenizer_name = "unified_tokenizer"

    model_data_bytes = redis_client.get(f"model:{model_name}")
    tokenizer_data_bytes = redis_client.get(f"tokenizer:{tokenizer_name}")

    if model_data_bytes:
        model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3)
        model.load_state_dict(torch.load(model_data_bytes))
    else:
        model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3)
    models[model_name] = model

    if tokenizer_data_bytes:
        tokenizer = AutoTokenizer.from_pretrained("gpt2")
        tokenizer.add_tokens(json.loads(tokenizer_data_bytes))
    else:
        tokenizer = AutoTokenizer.from_pretrained("gpt2")
    tokenizers[tokenizer_name] = tokenizer

    unified_model = UnifiedModel(list(models.values()))
    unified_model.to(torch.device("cpu"))

    if data.get("train"):
        user_data = data.get("user_data", [])
        if not user_data:
            user_data = [
                {"text": "Hola", "label": 1},
                {"text": "Necesito ayuda", "label": 2},
                {"text": "No entiendo", "label": 0}
            ]

        redis_client.rpush("training_queue", json.dumps({
            "tokenizers": {tokenizer_name: tokenizer.get_vocab()},
            "data": user_data
        }))

        return {"message": "Training data received. Model will be updated asynchronously."}

    elif data.get("message"):
        user_id = data.get("user_id")
        text = data['message']
        language = data.get("language", default_language)

        if user_id not in conversation_history:
            conversation_history[user_id] = []
        conversation_history[user_id].append(text)

        contextualized_text = " ".join(conversation_history[user_id][-3:])

        tokenized_inputs = [tokenizers[name](contextualized_text, return_tensors="pt") for name in tokenizers.keys()]
        input_ids = [tokens['input_ids'] for tokens in tokenized_inputs]
        attention_mask = [tokens['attention_mask'] for tokens in tokenized_inputs]

        with torch.no_grad():
            logits = unified_model(input_ids=input_ids, attention_mask=attention_mask)
            predicted_class = torch.argmax(logits, dim=-1).item()

        response = chatbot_service.get_response(user_id, contextualized_text, language)

        redis_client.rpush("training_queue", json.dumps({
            "tokenizers": {tokenizer_name: tokenizer.get_vocab()},
            "data": [{"text": contextualized_text, "label": predicted_class}]
        }))

        return {"answer": response}

    else:
        raise HTTPException(status_code=400, detail="Request must contain 'train' or 'message'.")

def get_chatbot_response(user_id, question, predicted_class, language):
    if user_id not in conversation_history:
        conversation_history[user_id] = []
    conversation_history[user_id].append(question)

    return chatbot_service.get_response(user_id, question, language)

@app.get("/")
async def get_home():
    user_id = str(uuid.uuid4())
    html_code = f"""
    <!DOCTYPE html>
    <html>
    <head>
        <meta charset="UTF-8">
        <title>Chatbot</title>
        <style>
            body {{
                font-family: 'Arial', sans-serif;
                background-color: #f4f4f9;
                margin: 0;
                padding: 0;
                display: flex;
                align-items: center;
                justify-content: center;
                min-height: 100vh;
            }}

            .container {{
                background-color: #fff;
                border-radius: 10px;
                box-shadow: 0 2px 5px rgba(0, 0, 0, 0.1);
                overflow: hidden;
                width: 400px;
                max-width: 90%;
            }}

            h1 {{
                color: #333;
                text-align: center;
                padding: 20px;
                margin: 0;
                background-color: #f8f9fa;
                border-bottom: 1px solid #eee;
            }}

            #chatbox {{
                height: 400px;
                padding: 20px;
                overflow-y: auto;
            }}

            .message {{
                margin-bottom: 15px;
                padding: 10px;
                border-radius: 5px;
                max-width: 70%;
                animation: slide-in 0.3s ease-out;
            }}

            .user-message {{
                text-align: right;
                background-color: #eee;
                margin-left: 30%;
            }}

            .bot-message {{
                text-align: left;
                background-color: #ccf5ff;
                margin-right: 30%;
            }}

            #input-area {{
                display: flex;
                padding: 10px;
                background-color: #f8f9fa;
                border-top: 1px solid #eee;
            }}

            #message-input {{
                flex: 1;
                padding: 10px;
                border: 1px solid #ccc;
                border-radius: 5px;
                margin-right: 10px;
            }}

            #send-button {{
                padding: 10px 15px;
                background-color: #28a745;
                color: white;
                border: none;
                cursor: pointer;
                border-radius: 5px;
                transition: background-color 0.3s ease;
            }}

            #send-button:hover {{
                background-color: #218838;
            }}

            @keyframes slide-in {{
                from {{
                    transform: translateX(-100%);
                    opacity: 0;
                }}
                to {{
                    transform: translateX(0);
                    opacity: 1;
                }}
            }}
        </style>
    </head>
    <body>
        <div class="container">
            <h1>Chatbot</h1>
            <div id="chatbox"></div>
            <div id="input-area">
                <input type="hidden" id="user-id" value="{user_id}">
                <input type="text" id="message-input" placeholder="Escribe tu mensaje...">
                <button id="send-button">Enviar</button>
            </div>
        </div>
        <script>
            const chatbox = document.getElementById('chatbox');
            const messageInput = document.getElementById('message-input');
            const sendButton = document.getElementById('send-button');
            const userId = document.getElementById('user-id').value;

            sendButton.addEventListener('click', sendMessage);

            function sendMessage() {{
                const message = messageInput.value;
                if (message.trim() === '') return;

                appendMessage('user', message);
                messageInput.value = '';

                fetch('/process', {{
                    method: 'POST',
                    headers: {{
                        'Content-Type': 'application/json'
                    }},
                    body: JSON.stringify({{ message: message, user_id: userId, language: 'es' }})
                }})
                .then(response => response.json())
                .then(data => {{
                    appendMessage('bot', data.answer);
                }});
            }}

            function appendMessage(sender, message) {{
                const messageElement = document.createElement('div');
                messageElement.classList.add('message', `${{sender}}-message`);
                messageElement.textContent = message;
                chatbox.appendChild(messageElement);
                chatbox.scrollTop = chatbox.scrollHeight;
            }}
        </script>
    </body>
    </html>
    """
    return HTMLResponse(content=html_code)

def push_to_redis(models, tokenizers, redis_client, model_name, tokenizer_name):
    for model_name, model in models.items():
        torch.save(model.state_dict(), model_name)
        with open(model_name, "rb") as f:
            redis_client.set(f"model:{model_name}", f.read())

    for tokenizer_name, tokenizer in tokenizers.items():
        tokens = tokenizer.get_vocab()
        redis_client.set(f"tokenizer:{tokenizer_name}", json.dumps(tokens))

def continuous_training():
    redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)

    while True:
        try:
            data = redis_client.lpop("training_queue")
            if data:
                data = json.loads(data)
                unified_model = UnifiedModel.load_model_from_redis(redis_client)
                unified_model.train()

                train_dataset = SyntheticDataset(data["tokenizers"], data["data"])
                train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)

                optimizer = AdamW(unified_model.parameters(), lr=5e-5)

                for epoch in range(10):
                    for batch in train_loader:
                        input_ids = [batch[f"input_ids_{name}"].to("cpu") for name in data["tokenizers"].keys()]
                        attention_mask = [batch[f"attention_mask_{name}"].to("cpu") for name in data["tokenizers"].keys()]
                        labels = batch["labels"].to("cpu")
                        outputs = unified_model(input_ids=input_ids, attention_mask=attention_mask)
                        loss = nn.CrossEntropyLoss()(outputs, labels)
                        loss.backward()
                        optimizer.step()
                        optimizer.zero_grad()

                        print(f"Epoch {epoch}, Loss {loss.item()}")

                push_to_redis(
                    {"response_model": unified_model},
                    {"response_tokenizer": tokenizer},
                    redis_client,
                    "response_model",
                    "response_tokenizer",
                )
            time.sleep(10)
        except Exception as e:
            print(f"Error in continuous training: {e}")
            time.sleep(5)

if __name__ == "__main__":
    training_process = multiprocessing.Process(target=continuous_training)
    training_process.start()

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