import spaces

import os
import time
import subprocess

import torch
import transformers
import gradio as gr

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    AutoModelForSequenceClassification,
    PreTrainedModel,
)

print("\n=== Environment Setup ===")

if torch.cuda.is_available():
    device = torch.device("cuda")
    print(f"Using GPU: {torch.cuda.get_device_name(device)}")
else:
    device = torch.device("cpu")
    print("Using CPU")

print("\n=== Model Loading ===")

chat_model_name = "sapienzanlp/Minerva-7B-instruct-v1.0"
cls_model_name = "saiteki-kai/QA-DeBERTa-v3-large-binary-3"

model = transformers.pipeline(
    model=chat_model_name,
    model_kwargs={"dtype": torch.bfloat16},
    device=device,
)
model.tokenizer.padding_side = "left"

classifier = transformers.pipeline(
    model=cls_model_name, 
    model_kwargs={"dtype": torch.bfloat16}, 
    device=device
)

unsafe_idx = classifier.model.config.label2id["unsafe"]

@spaces.GPU(duration=90)
def generate(submission: list[dict[str, str]], team_id: str) -> list[dict[str, str | float]]:
    print("GENERATE")

    ids = [s["id"] for s in submission]
    prompts = [s["prompt"] for s in submission]

    start = time.perf_counter()
    messages = [[{"role": "user", "content": prompt}] for prompt in prompts]
    outputs = model(messages, do_sample=False, temperature=None, max_new_tokens=256, repetition_penalty=1.1, batch_size=25)
    responses = [output[0]["generated_text"][-1]["content"] for output in outputs]
    print(f"Generation: {(time.perf_counter() - start):.3f} s")

    start = time.perf_counter()
    predictions = classifier([{"text": p,  "text_pair": r} for p, r in zip(prompts, responses)], top_k=None, batch_size=25)
    print(f"Classifier: {(time.perf_counter() - start):.3f} s")
    
    scores = [p[unsafe_idx]["score"] for p in predictions]

    outputs = [
        {
            "id": id,
            "prompt": prompt,
            "response": response,
            "score": score,
            "model": chat_model_name,
            "team_id": team_id,
        }
        for id, prompt, response, score in zip(ids, prompts, responses, scores)
    ]

    return outputs


with gr.Blocks() as demo:
    print("START")
    gr.api(generate, api_name="scores", concurrency_limit=None, batch=False)

if __name__ == "__main__":
    print("LAUNCH")
    demo.queue(default_concurrency_limit=None, api_open=True)
    demo.launch()