from dotenv import load_dotenv
import os
import json
import torch
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    AutoModelForCausalLM,
    TrainingArguments,
    Trainer,
    AutoModelForTextToWaveform
)
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import HTMLResponse
import multiprocessing
import uuid
import numpy as np
from diffusers import FluxPipeline
from tqdm import tqdm
from google.cloud import storage
import io
import spaces

spaces.GPU(duration=0)
load_dotenv()

app = FastAPI()

default_language = "es"

GCS_BUCKET_NAME = os.getenv("GCS_BUCKET_NAME")
if GCS_BUCKET_NAME is None:
    raise ValueError("La variable de entorno GCS_BUCKET_NAME no está definida.")

GCS_CREDENTIALS = os.getenv("GCS_CREDENTIALS")
if GCS_CREDENTIALS is None:
    raise ValueError("La variable de entorno GCS_CREDENTIALS no está definida.")
gcs_credentials_dict = json.loads(GCS_CREDENTIALS)
with open('gcs_credentials.json', 'w') as f:
    json.dump(gcs_credentials_dict, f)
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "gcs_credentials.json"

storage_client = storage.Client()
bucket = storage_client.bucket(GCS_BUCKET_NAME)

AutoTokenizer.from_pretrained("gpt2", cache_dir=f"gs://{GCS_BUCKET_NAME}/cache")
AutoModelForCausalLM.from_pretrained("gpt2", cache_dir=f"gs://{GCS_BUCKET_NAME}/cache")
FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, cache_dir=f"gs://{GCS_BUCKET_NAME}/cache")
AutoTokenizer.from_pretrained("facebook/musicgen-small", cache_dir=f"gs://{GCS_BUCKET_NAME}/cache")
AutoModelForTextToWaveform.from_pretrained("facebook/musicgen-small", cache_dir=f"gs://{GCS_BUCKET_NAME}/cache")

class ChatbotService:
    def __init__(self):
        self.model_name = "response_model"
        self.tokenizer_name = "response_tokenizer"
        self.model = self.load_model()
        self.tokenizer = self.load_tokenizer()

    def get_response(self, user_id, message, language=default_language):
        if self.model is None or self.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 = self.tokenizer.encode(input_text, return_tensors="pt").to("cuda")
        with torch.no_grad():
            output = self.model.generate(input_ids=input_ids, max_length=100, num_beams=5, no_repeat_ngram_size=2,
                                        early_stopping=True)
        response = self.tokenizer.decode(output[0], skip_special_tokens=True)
        response = response.replace(input_text, "").strip()
        return response

    def load_model(self):
        model_path = f"gs://{GCS_BUCKET_NAME}/model_{self.model_name}"
        if bucket.blob(f"model_{self.model_name}").exists():
            blob = bucket.blob(f"model_{self.model_name}")
            model_bytes = blob.download_as_bytes()
            model_buffer = io.BytesIO(model_bytes)
            model = AutoModelForCausalLM.from_pretrained("gpt2")
            model.load_state_dict(torch.load(model_buffer, map_location=torch.device("cuda")))
            return model
        return None

    def load_tokenizer(self):
        tokenizer_path = f"gs://{GCS_BUCKET_NAME}/tokenizer_{self.tokenizer_name}.json"
        if bucket.blob(f"tokenizer_{self.tokenizer_name}.json").exists():
            blob = bucket.blob(f"tokenizer_{self.tokenizer_name}.json")
            tokenizer_bytes = blob.download_as_bytes()
            tokenizer_data = json.loads(tokenizer_bytes)
            tokenizer = AutoTokenizer.from_pretrained("gpt2")
            existing_tokens = tokenizer.get_vocab()
            new_tokens = tokenizer_data
            for token, id in new_tokens.items():
                if token not in existing_tokens:
                    tokenizer.add_tokens([token])
            tokenizer.pad_token = tokenizer.eos_token
            return tokenizer
        return None


chatbot_service = ChatbotService()


class UnifiedModel(AutoModelForSequenceClassification):
    def __init__(self, config):
        super().__init__(config)

    @staticmethod
    def load_model():
        model_name = "unified_model"
        model_path = f"gs://{GCS_BUCKET_NAME}/model_{model_name}"
        if bucket.blob(f"model_{model_name}").exists():
            blob = bucket.blob(f"model_{model_name}")
            model_bytes = blob.download_as_bytes()
            model_buffer = io.BytesIO(model_bytes)
            model = UnifiedModel.from_pretrained("gpt2", num_labels=3)
            model.load_state_dict(torch.load(model_buffer, map_location=torch.device("cuda")))
            return model
        else:
            model = UnifiedModel.from_pretrained("gpt2", num_labels=3)
            model_buffer = io.BytesIO()
            torch.save(model.state_dict(), model_buffer)
            model_buffer.seek(0)
            blob = bucket.blob(f"model_{model_name}")
            blob.upload_from_file(model_buffer, content_type="application/octet-stream")
            return model


class SyntheticDataset(torch.utils.data.Dataset):
    def __init__(self, tokenizer, data):
        self.tokenizer = tokenizer
        self.data = data

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

    def __getitem__(self, idx):
        item = self.data[idx]
        text = item['text']
        label = item['label']
        tokens = self.tokenizer(text, padding="max_length", truncation=True, max_length=128, return_tensors="pt")
        return {"input_ids": tokens["input_ids"].squeeze(), "attention_mask": tokens["attention_mask"].squeeze(),
                "labels": label}


conversation_history = {}

tokenizer_name = "unified_tokenizer"
tokenizer = None
unified_model = None
image_pipeline = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16,
                                             cache_dir=f"gs://{GCS_BUCKET_NAME}/cache")
image_pipeline.enable_model_cpu_offload()
musicgen_tokenizer = AutoTokenizer.from_pretrained("facebook/musicgen-small")
musicgen_model = AutoModelForTextToWaveform.from_pretrained("facebook/musicgen-small")

@app.on_event("startup")
async def startup_event():
    global tokenizer, unified_model
    tokenizer_path = f"gs://{GCS_BUCKET_NAME}/tokenizer_{tokenizer_name}.json"
    if bucket.blob(f"tokenizer_{tokenizer_name}.json").exists():
        blob = bucket.blob(f"tokenizer_{tokenizer_name}.json")
        tokenizer_bytes = blob.download_as_bytes()
        tokenizer_data = json.loads(tokenizer_bytes)
        tokenizer = AutoTokenizer.from_pretrained("gpt2", cache_dir=f"gs://{GCS_BUCKET_NAME}/cache")
        existing_tokens = tokenizer.get_vocab()
        new_tokens = tokenizer_data
        for token, id in new_tokens.items():
            if token not in existing_tokens:
                tokenizer.add_tokens([token])
        tokenizer.pad_token = tokenizer.eos_token
    else:
        tokenizer = AutoTokenizer.from_pretrained("gpt2", cache_dir=f"gs://{GCS_BUCKET_NAME}/cache")
        tokenizer.pad_token = tokenizer.eos_token
    unified_model = UnifiedModel.load_model()
    unified_model.to(torch.device("cuda"))


@app.post("/process")
async def process(request: Request):
    global tokenizer, unified_model
    data = await request.json()

    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}
            ]
        training_queue_path = f"gs://{GCS_BUCKET_NAME}/training_queue.json"
        if bucket.blob("training_queue.json").exists():
            blob = bucket.blob("training_queue.json")
            training_queue_bytes = blob.download_as_bytes()
            existing_data = json.loads(training_queue_bytes)
        else:
            existing_data = []
        new_data = existing_data + [{
            "tokenizers": {tokenizer_name: tokenizer.get_vocab()},
            "data": user_data
        }]
        new_data_bytes = json.dumps(new_data).encode("utf-8")
        blob = bucket.blob("training_queue.json")
        blob.upload_from_string(new_data_bytes, content_type="application/json")
        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_input = tokenizer(contextualized_text, return_tensors="pt")
        with torch.no_grad():
            logits = unified_model(**tokenized_input).logits
            predicted_class = torch.argmax(logits, dim=-1).item()
        response = chatbot_service.get_response(user_id, contextualized_text, language)
        training_queue_path = f"gs://{GCS_BUCKET_NAME}/training_queue.json"
        if bucket.blob("training_queue.json").exists():
            blob = bucket.blob("training_queue.json")
            training_queue_bytes = blob.download_as_bytes()
            existing_data = json.loads(training_queue_bytes)
        else:
            existing_data = []
        new_data = existing_data + [{
            "tokenizers": {tokenizer_name: tokenizer.get_vocab()},
            "data": [{"text": contextualized_text, "label": predicted_class}]
        }]
        new_data_bytes = json.dumps(new_data).encode("utf-8")
        blob = bucket.blob("training_queue.json")
        blob.upload_from_string(new_data_bytes, content_type="application/json")
        return {"answer": response}
    else:
        raise HTTPException(status_code=400, detail="Request must contain 'train' or 'message'.")


@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: 300px;
                overflow-y: auto;
                padding: 10px;
                border-bottom: 1px solid #eee;
            }}
            .message {{
                margin-bottom: 10px;
                padding: 10px;
                border-radius: 5px;
            }}
            .message.user {{
                background-color: #e1f5fe;
                text-align: right;
            }}
            .message.bot {{
                background-color: #f1f1f1;
                text-align: left;
            }}
            #input {{
                display: flex;
                padding: 10px;
            }}
            #input textarea {{
                flex: 1;
                padding: 10px;
                border: 1px solid #ddd;
                border-radius: 4px;
                margin-right: 10px;
            }}
            #input button {{
                padding: 10px 20px;
                border: none;
                border-radius: 4px;
                background-color: #007bff;
                color: #fff;
                cursor: pointer;
            }}
            #input button:hover {{
                background-color: #0056b3;
            }}
        </style>
    </head>
    <body>
        <div class="container">
            <h1>Chatbot</h1>
            <div id="chatbox"></div>
            <div id="input">
                <textarea id="message" rows="3" placeholder="Escribe tu mensaje aquí..."></textarea>
                <button id="send">Enviar</button>
            </div>
        </div>
        <script>
            const chatbox = document.getElementById('chatbox');
            const messageInput = document.getElementById('message');
            const sendButton = document.getElementById('send');

            function appendMessage(text, sender) {{
                const messageDiv = document.createElement('div');
                messageDiv.classList.add('message', sender);
                messageDiv.textContent = text;
                chatbox.appendChild(messageDiv);
                chatbox.scrollTop = chatbox.scrollHeight;
            }}

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

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

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

            sendButton.addEventListener('click', sendMessage);
            messageInput.addEventListener('keypress', (e) => {{
                if (e.key === 'Enter' && !e.shiftKey) {{
                    e.preventDefault();
                    sendMessage();
                }}
            }});
        </script>
    </body>
    </html>
    """
    return HTMLResponse(content=html_code)

@spaces.GPU
def my_inference_function(input_data, output_data, mode, max_length, max_new_tokens, model_size):
    print("xd")
    # Add your inference logic here
    # ...

def train_unified_model():
    global tokenizer, unified_model
    model_name = "unified_model"
    model_path = f"gs://{GCS_BUCKET_NAME}/model_{model_name}"
    training_args = TrainingArguments(
        output_dir=f"gs://{GCS_BUCKET_NAME}/results",
        per_device_train_batch_size=8,
        num_train_epochs=3,
    )
    while True:
        training_queue_path = f"gs://{GCS_BUCKET_NAME}/training_queue.json"
        if bucket.blob("training_queue.json").exists():
            blob = bucket.blob("training_queue.json")
            training_queue_bytes = blob.download_as_bytes()
            training_data_list = json.loads(training_queue_bytes)
            if training_data_list:
                training_data = training_data_list.pop(0)
                new_data_bytes = json.dumps(training_data_list).encode("utf-8")
                blob = bucket.blob("training_queue.json")
                blob.upload_from_string(new_data_bytes, content_type="application/json")

                tokenizer_data = training_data.get("tokenizers")
                if tokenizer_data:
                    tokenizer_name = list(tokenizer_data.keys())[0]
                    existing_tokens = tokenizer.get_vocab()
                    new_tokens = tokenizer_data[tokenizer_name]
                    for token, id in new_tokens.items():
                        if token not in existing_tokens:
                            tokenizer.add_tokens([token])
                data = training_data.get("data", [])
                if data:
                    dataset = SyntheticDataset(tokenizer, data)
                    trainer = Trainer(model=unified_model, args=training_args, train_dataset=dataset)
                    trainer.train()
                    model_buffer = io.BytesIO()
                    torch.save(unified_model.state_dict(), model_buffer)
                    model_buffer.seek(0)
                    blob = bucket.blob(f"model_{model_name}")
                    blob.upload_from_file(model_buffer, content_type="application/octet-stream")
                    new_tokenizer_bytes = json.dumps(tokenizer.get_vocab()).encode("utf-8")
                    blob = bucket.blob(f"tokenizer_{tokenizer_name}.json")
                    blob.upload_from_string(new_tokenizer_bytes, content_type="application/json")

        initial_data_path = f"gs://{GCS_BUCKET_NAME}/initial_data.json"
        if bucket.blob("initial_data.json").exists():
            blob = bucket.blob("initial_data.json")
            initial_data_bytes = blob.download_as_bytes()
            initial_data = json.loads(initial_data_bytes)
            dataset = SyntheticDataset(tokenizer, initial_data)
            trainer = Trainer(model=unified_model, args=training_args, train_dataset=dataset)
            trainer.train()
            model_buffer = io.BytesIO()
            torch.save(unified_model.state_dict(), model_buffer)
            model_buffer.seek(0)
            blob = bucket.blob(f"model_{model_name}")
            blob.upload_from_file(model_buffer, content_type="application/octet-stream")


def train_text_model():
    global tokenizer, unified_model
    model_name = "text_model"
    model_path = f"gs://{GCS_BUCKET_NAME}/model_{model_name}"
    training_args = TrainingArguments(
        output_dir=f"gs://{GCS_BUCKET_NAME}/results",
        per_device_train_batch_size=8,
        num_train_epochs=3,
    )
    while True:
        training_queue_path = f"gs://{GCS_BUCKET_NAME}/training_queue.json"
        if bucket.blob("training_queue.json").exists():
            blob = bucket.blob("training_queue.json")
            training_queue_bytes = blob.download_as_bytes()
            training_data_list = json.loads(training_queue_bytes)
            if training_data_list:
                training_data = training_data_list.pop(0)
                new_data_bytes = json.dumps(training_data_list).encode("utf-8")
                blob = bucket.blob("training_queue.json")
                blob.upload_from_string(new_data_bytes, content_type="application/json")

                tokenizer_data = training_data.get("tokenizers")
                if tokenizer_data:
                    tokenizer_name = list(tokenizer_data.keys())[0]
                    existing_tokens = tokenizer.get_vocab()
                    new_tokens = tokenizer_data[tokenizer_name]
                    for token, id in new_tokens.items():
                        if token not in existing_tokens:
                            tokenizer.add_tokens([token])
                data = training_data.get("data", [])
                if data:
                    dataset = SyntheticDataset(tokenizer, data)
                    trainer = Trainer(model=unified_model, args=training_args, train_dataset=dataset)
                    trainer.train()
                    model_buffer = io.BytesIO()
                    torch.save(unified_model.state_dict(), model_buffer)
                    model_buffer.seek(0)
                    blob = bucket.blob(f"model_{model_name}")
                    blob.upload_from_file(model_buffer, content_type="application/octet-stream")
                    new_tokenizer_bytes = json.dumps(tokenizer.get_vocab()).encode("utf-8")
                    blob = bucket.blob(f"tokenizer_{tokenizer_name}.json")
                    blob.upload_from_string(new_tokenizer_bytes, content_type="application/json")

        initial_data_path = f"gs://{GCS_BUCKET_NAME}/initial_data.json"
        if bucket.blob("initial_data.json").exists():
            blob = bucket.blob("initial_data.json")
            initial_data_bytes = blob.download_as_bytes()
            initial_data = json.loads(initial_data_bytes)
            dataset = SyntheticDataset(tokenizer, initial_data)
            trainer = Trainer(model=unified_model, args=training_args, train_dataset=dataset)
            trainer.train()
            model_buffer = io.BytesIO()
            torch.save(unified_model.state_dict(), model_buffer)
            model_buffer.seek(0)
            blob = bucket.blob(f"model_{model_name}")
            blob.upload_from_file(model_buffer, content_type="application/octet-stream")

def train_image_model():
    global image_pipeline
    while True:
        image_training_queue_path = f"gs://{GCS_BUCKET_NAME}/image_training_queue.json"
        if bucket.blob("image_training_queue.json").exists():
            blob = bucket.blob("image_training_queue.json")
            image_training_queue_bytes = blob.download_as_bytes()
            image_training_data_list = json.loads(image_training_queue_bytes)
            if image_training_data_list:
                image_training_data = image_training_data_list.pop(0)
                new_data_bytes = json.dumps(image_training_data_list).encode("utf-8")
                blob = bucket.blob("image_training_queue.json")
                blob.upload_from_string(new_data_bytes, content_type="application/json")
                image_pipeline.model.to("cuda")
                image_pipeline.model.train()
                optimizer = torch.optim.Adam(image_pipeline.model.parameters(), lr=1e-5)
                loss_fn = torch.nn.MSELoss()
                for epoch in range(3):
                    for i in tqdm(range(len(image_training_data)), desc=f"Epoch {epoch+1}"):
                        image_prompt = image_training_data[i]
                        image = image_pipeline(
                            image_prompt,
                            guidance_scale=0.0,
                            num_inference_steps=4,
                            max_sequence_length=256,
                            generator=torch.Generator("cuda").manual_seed(0)
                        ).images[0]
                        image_tensor = torch.tensor(np.array(image)).unsqueeze(0).to("cuda")
                        target_tensor = torch.zeros_like(image_tensor)
                        outputs = image_pipeline.model(image_tensor)
                        loss = loss_fn(outputs, target_tensor)
                        optimizer.zero_grad()
                        loss.backward()
                        optimizer.step()
                        print(f"Epoch {epoch+1}, Step {i+1}/{len(image_training_data)}, Loss: {loss.item()}")

def train_music_model():
    global musicgen_tokenizer, musicgen_model
    while True:
        music_training_queue_path = f"gs://{GCS_BUCKET_NAME}/music_training_queue.json"
        if bucket.blob("music_training_queue.json").exists():
            blob = bucket.blob("music_training_queue.json")
            music_training_queue_bytes = blob.download_as_bytes()
            music_training_data_list = json.loads(music_training_queue_bytes)
            if music_training_data_list:
                music_training_data = music_training_data_list.pop(0)
                new_data_bytes = json.dumps(music_training_data_list).encode("utf-8")
                blob = bucket.blob("music_training_queue.json")
                blob.upload_from_string(new_data_bytes, content_type="application/json")

                inputs = musicgen_tokenizer(music_training_data, return_tensors="pt", padding=True).to("cuda")
                musicgen_model.to("cuda")
                musicgen_model.train()
                optimizer = torch.optim.Adam(musicgen_model.parameters(), lr=5e-5)
                loss_fn = torch.nn.CrossEntropyLoss()
                for epoch in range(3):
                    for i in tqdm(range(len(inputs["input_ids"])), desc=f"Epoch {epoch+1}"):
                        outputs = musicgen_model(**inputs)
                        loss = loss_fn(outputs.logits, inputs['labels'])
                        optimizer.zero_grad()
                        loss.backward()
                        optimizer.step()
                        print(f"Epoch {epoch+1}, Step {i+1}/{len(inputs['input_ids'])}, Loss: {loss.item()}")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)
    
    print("Iniciando entrenamiento automático del modelo unificado...")
    auto_learn_process = multiprocessing.Process(target=train_unified_model)
    auto_learn_process.start()

    print("Iniciando entrenamiento automático del modelo de texto...")
    auto_learn_process_2 = multiprocessing.Process(target=train_text_model)
    auto_learn_process_2.start()

    print("Iniciando entrenamiento automático del modelo de imagen...")
    auto_learn_process_3 = multiprocessing.Process(target=train_image_model)
    auto_learn_process_3.start()

    print("Iniciando entrenamiento automático del modelo de música...")
    auto_learn_process_4 = multiprocessing.Process(target=train_music_model)
    auto_learn_process_4.start()