import os
import gradio as gr
import pandas as pd
from datasets import Dataset
from transformers import (
    T5Tokenizer,
    T5ForConditionalGeneration,
    Trainer,
    TrainingArguments
)
from huggingface_hub import login
from spaces import GPU  # Required for ZeroGPU Spaces

import transformers
print("🔥 Transformers version:", transformers.__version__)

import torch
device = torch.device("cpu")

login(token=os.getenv("HF_TOKEN"))  # Using environment variable
# Disable CUDA and set environment variables for Spaces with Stateless GPU
os.environ["CUDA_VISIBLE_DEVICES"] = ""  # Ensure no GPUs are used
os.environ["ACCELERATE_DISABLE"] = "true"  # Disable Accelerate for safety

model_name = "madankn/xindus_t5base"

# Load tokenizer and model once
tokenizer = T5Tokenizer.from_pretrained("t5-base")
model = T5ForConditionalGeneration.from_pretrained("t5-base")
model.to("cpu")  # Ensure model stays on CPU in main process

def preprocess(example):
    inputs = tokenizer(
        example["text"],
        padding="max_length",
        truncation=True,
        max_length=512,
    )
    labels = tokenizer(
        example["summary"],
        padding="max_length",
        truncation=True,
        max_length=128,
    )
    return {
        "input_ids": inputs["input_ids"],
        "attention_mask": inputs["attention_mask"],
        "labels": labels["input_ids"]
    }

# Load fine-tuned model for inference
def load_model():
    global tokenizer, model
    if tokenizer is None or model is None:
        tokenizer = T5Tokenizer.from_pretrained("./fine_tuned_t5")
        model = T5ForConditionalGeneration.from_pretrained("./fine_tuned_t5")
        model.to("cpu")
    return model, tokenizer

def train_model():
    # 🔁 Reload model/tokenizer inside training function to avoid stale GPU bindings
    tokenizer = T5Tokenizer.from_pretrained("t5-base")
    model = T5ForConditionalGeneration.from_pretrained("t5-base").to(device)

    df = pd.read_csv("xindus_dataset.csv")
    df = df.rename(columns={"text_column_name": "text", "summary_column_name": "summary"})
    dataset = Dataset.from_pandas(df).train_test_split(test_size=0.1)

    tokenized_datasets = dataset.map(
        preprocess,
        batched=True,
        remove_columns=dataset["train"].column_names
    )

    training_args = TrainingArguments(
        output_dir="./results",
        logging_dir="./logs",
        logging_steps=50,
        save_steps=200,
        num_train_epochs=1,
        per_device_train_batch_size=2,
        per_device_eval_batch_size=2,
        weight_decay=0.01,
        learning_rate=2e-5,
        save_total_limit=1,
        push_to_hub=True,
        hub_model_id=model_name,
        hub_strategy="every_save",
        no_cuda=True  # 🧠 Critical: disable all CUDA use
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_datasets["train"],
        eval_dataset=tokenized_datasets["test"],
        tokenizer=tokenizer,
    )

    trainer.train()
    model.save_pretrained("./fine_tuned_t5")
    tokenizer.save_pretrained("./fine_tuned_t5")
    model.push_to_hub(model_name)
    tokenizer.push_to_hub(model_name)

    return f"✅ Training complete and pushed to: https://huggingface.co/{model_name}"

# Summarize function using GPU
@GPU
def summarize(text):
    model, tokenizer = load_model()
    input_ids = tokenizer("summarize: " + text, return_tensors="pt", truncation=True).input_ids
    output_ids = model.generate(input_ids, max_length=50)
    return tokenizer.decode(output_ids[0], skip_special_tokens=True)

# Gradio UI
train_button = gr.Interface(fn=train_model, inputs=[], outputs="text", title="Train T5 on Xindus Data")
summarize_interface = gr.Interface(fn=summarize, inputs="text", outputs="text", title="Summarize with Fine-Tuned T5")

# Combine interfaces
def combined_interface():
    with gr.Blocks() as demo:
        with gr.Tab("Training"):
            train_button.render()
        with gr.Tab("Summarization"):
            summarize_interface.render()
    demo.launch()

# Launch app
combined_interface()