from dotenv import load_dotenv
import os
import json
import requests
import redis
from transformers import (
    AutoTokenizer,
    AutoModel,
    TrainingArguments,
)
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 pydantic import BaseModel
from typing import List, Dict
from fastapi.responses import HTMLResponse
import multiprocessing
import time

load_dotenv()

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

app = FastAPI()

class UnifiedModel(nn.Module):
    def __init__(self, models):
        super(UnifiedModel, self).__init__()
        self.models = nn.ModuleList(models)
        self.classifier = nn.Linear(sum([model.config.hidden_size for model in models]), 2)

    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.last_hidden_state[:, 0, :])
        concatenated_hidden_states = torch.cat(hidden_states, dim=-1)
        logits = self.classifier(concatenated_hidden_states)
        return logits

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["label"] = torch.tensor(label)
        return tokenized

@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_data = json.loads(model_data_bytes)
        model = AutoModel.from_pretrained("gpt2")
        model.load_state_dict(torch.load(model_data))
    else:
        model = AutoModel.from_pretrained("gpt2")
    models[model_name] = model

    if tokenizer_data_bytes:
        tokenizer_data = json.loads(tokenizer_data_bytes)
        tokenizer = AutoTokenizer.from_pretrained("gpt2")
        tokenizer.add_tokens(tokenizer_data)
    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": "Sample text for automatic training.", "label": 0}]
        
        train_dataset = SyntheticDataset(tokenizers, user_data)
        train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
        
        training_args = TrainingArguments(
            output_dir="memory",
            evaluation_strategy="epoch",
            learning_rate=5e-5,
            per_device_train_batch_size=8,
            per_device_eval_batch_size=8,
            num_train_epochs=10,
            weight_decay=0.01,
            logging_steps=10,
            optim="adamw_hf"
        )
        
        optimizer = AdamW(unified_model.parameters(), lr=training_args.learning_rate)
        unified_model.train()
        
        for epoch in range(training_args.num_train_epochs):
            for batch in train_loader:
                input_ids = [batch[f"input_ids_{name}"].to("cpu") for name in tokenizers.keys()]
                attention_mask = [batch[f"attention_mask_{name}"].to("cpu") for name in tokenizers.keys()]
                labels = batch["label"].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()}")
        
        print("Training complete.")

        push_to_redis(models, tokenizers, redis_client, model_name, tokenizer_name)
        return {"message": "Model trained and updated in Redis."}
    
    elif data.get("predict"):
        text = data['text']
        tokenized_inputs = [tokenizers[name](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()
        
        return {"prediction": predicted_class}
    
    else:
        raise HTTPException(status_code=400, detail="Request must contain 'train' or 'predict'.")

@app.post("/external_answer")
async def external_answer(request: Request):
    data = await request.json()
    redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)
    
    question = data.get('question')
    if not question:
        raise HTTPException(status_code=400, detail="Question is required.")

    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_data = json.loads(model_data_bytes)
        model = AutoModel.from_pretrained("gpt2")
        model.load_state_dict(torch.load(model_data))
    else:
        model = AutoModel.from_pretrained("gpt2")

    if tokenizer_data_bytes:
        tokenizer_data = json.loads(tokenizer_data_bytes)
        tokenizer = AutoTokenizer.from_pretrained("gpt2")
        tokenizer.add_tokens(tokenizer_data)
    else:
        tokenizer = AutoTokenizer.from_pretrained("gpt2")

    unified_model = UnifiedModel([model])
    unified_model.to(torch.device("cpu"))
    
    tokenized_input = tokenizer(question, return_tensors="pt")
    input_ids = tokenized_input['input_ids']
    attention_mask = tokenized_input['attention_mask']
    
    with torch.no_grad():
        logits = unified_model(input_ids=input_ids, attention_mask=attention_mask)
        predicted_class = torch.argmax(logits, dim=-1).item()
        response = {"answer": f"Response to '{question}' is class {predicted_class}"}
    
    extreme_training_data = [{"text": question, "label": predicted_class}]
    train_dataset = SyntheticDataset({tokenizer_name: tokenizer}, extreme_training_data)
    train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
    
    training_args = TrainingArguments(
        output_dir="memory",
        evaluation_strategy="epoch",
        learning_rate=5e-5,
        per_device_train_batch_size=8,
        per_device_eval_batch_size=8,
        num_train_epochs=10,
        weight_decay=0.01,
        logging_steps=10,
        optim="adamw_hf"
    )
    
    optimizer = AdamW(unified_model.parameters(), lr=training_args.learning_rate)
    unified_model.train()
    
    for epoch in range(training_args.num_train_epochs):
        for batch in train_loader:
            input_ids = [batch[f"input_ids_{name}"].to("cpu") for name in [tokenizer_name]]
            attention_mask = [batch[f"attention_mask_{name}"].to("cpu") for name in [tokenizer_name]]
            labels = batch["label"].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()}")
    
    print("Extreme training complete.")
    push_to_redis({model_name: model}, {tokenizer_name: tokenizer}, redis_client, model_name, tokenizer_name)
    
    return response

@app.get("/")
async def get_home():
    html_code = """
    <!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;
            }
            .container {
                max-width: 1200px;
                margin: 0 auto;
                padding: 20px;
            }
            h1 {
                color: #333;
                text-align: center;
            }
            .grid-container {
                display: grid;
                grid-template-columns: repeat(auto-fill, minmax(300px, 1fr));
                gap: 10px;
                margin-top: 20px;
            }
            .grid-item {
                background: #fff;
                padding: 20px;
                border-radius: 8px;
                box-shadow: 0 0 10px rgba(0, 0, 0, 0.1);
                transition: transform 0.3s;
            }
            .grid-item:hover {
                transform: scale(1.05);
            }
            .question {
                font-weight: bold;
                color: #007bff;
            }
            .answer {
                margin-top: 10px;
                color: #333;
            }
            input[type="text"] {
                width: calc(100% - 22px);
                padding: 10px;
                margin: 0;
                border: 1px solid #ddd;
                border-radius: 4px;
            }
            button {
                padding: 10px 20px;
                background-color: #007bff;
                color: #fff;
                border: none;
                border-radius: 4px;
                cursor: pointer;
                margin-top: 10px;
            }
            button:hover {
                background-color: #0056b3;
            }
        </style>
        <script>
            async function sendMessage() {
                const question = document.getElementById('question').value;
                const responseElement = document.getElementById('response');
                
                const response = await fetch('/external_answer', {
                    method: 'POST',
                    headers: {
                        'Content-Type': 'application/json',
                    },
                    body: JSON.stringify({ question: question })
                });
                
                const data = await response.json();
                responseElement.innerText = "Response: " + data.answer;
                
                const gridContainer = document.getElementById('grid-container');
                const newItem = document.createElement('div');
                newItem.classList.add('grid-item');
                newItem.innerHTML = `<div class="question">${question}</div><div class="answer">${data.answer}</div>`;
                gridContainer.prepend(newItem);
            }
        </script>
    </head>
    <body>
        <div class="container">
            <h1>Chatbot</h1>
            <input type="text" id="question" placeholder="Ask me something...">
            <button onclick="sendMessage()">Send</button>
            <div id="response"></div>
            <div class="grid-container" id="grid-container"></div>
        </div>
    </body>
    </html>
    """
    return HTMLResponse(content=html_code)

def push_to_redis(models, tokenizers, redis_client, model_name, tokenizer_name):
    model_data = json.dumps(next(iter(models.values())).state_dict())
    redis_client.set(f"model:{model_name}", model_data)

    tokenizer_data = json.dumps(next(iter(tokenizers.values())).get_vocab())
    redis_client.set(f"tokenizer:{tokenizer_name}", tokenizer_data)

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

    model_name = "unified_model"
    tokenizer_name = "unified_tokenizer"

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

            if model_data_bytes and tokenizer_data_bytes:
                model_data = json.loads(model_data_bytes)
                model = AutoModel.from_pretrained("gpt2")
                model.load_state_dict(torch.load(model_data))

                tokenizer_data = json.loads(tokenizer_data_bytes)
                tokenizer = AutoTokenizer.from_pretrained("gpt2")
                tokenizer.add_tokens(tokenizer_data)

                unified_model = UnifiedModel([model])
                unified_model.to(torch.device("cpu"))

                train_data = [{"text": "Sample training text.", "label": 0}]
                train_dataset = SyntheticDataset({tokenizer_name: tokenizer}, train_data)
                train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)

                training_args = TrainingArguments(
                    output_dir="memory",
                    evaluation_strategy="epoch",
                    learning_rate=5e-5,
                    per_device_train_batch_size=8,
                    per_device_eval_batch_size=8,
                    num_train_epochs=10,
                    weight_decay=0.01,
                    logging_steps=10,
                    optim="adamw_hf"
                )

                optimizer = AdamW(unified_model.parameters(), lr=training_args.learning_rate)
                unified_model.train()

                for epoch in range(training_args.num_train_epochs):
                    for batch in train_loader:
                        input_ids = [batch[f"input_ids_{name}"].to("cpu") for name in [tokenizer_name]]
                        attention_mask = [batch[f"attention_mask_{name}"].to("cpu") for name in [tokenizer_name]]
                        labels = batch["label"].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()}")

                print("Training complete.")
                push_to_redis({model_name: model}, {tokenizer_name: tokenizer}, redis_client, model_name, tokenizer_name)
            else:
                print("No model or tokenizer found in Redis. Skipping training.")

            time.sleep(600)

        except Exception as e:
            print(f"An error occurred: {e}")
            time.sleep(60)

def start_server():
    import uvicorn
    cpu_cores = os.cpu_count() or 1
    num_workers = max(1, cpu_cores - 1)
    
    uvicorn.run(app, host="0.0.0.0", port=7860, timeout_keep_alive=0)

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