import gradio as gr
from peft import PeftModel
from transformers import RobertaTokenizer, T5ForConditionalGeneration, AutoTokenizer, AutoModelForCausalLM
import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Define model paths
model_id = {
    'CodeT5': "Salesforce/codet5-base",
    'StarCoder': "bigcode/starcoder2-3b",
    'CodeLlama': "codellama/CodeLlama-7b-Instruct-hf"
}

# Quantization Config
# quantization_config = BitsAndBytesConfig(
#     load_in_4bit=True,
#     bnb_4bit_quant_type="nf4",
#     bnb_4bit_compute_dtype=torch.bfloat16
# )

# Load CodeT5 models
codeT5_tokenizer = RobertaTokenizer.from_pretrained(model_id['CodeT5'])
base_codeT5 = T5ForConditionalGeneration.from_pretrained(model_id['CodeT5']).to(device)
base_codeT5.eval()
original_codeT5 = T5ForConditionalGeneration.from_pretrained(model_id['CodeT5']).to(device)
fine_tuned_codeT5 = PeftModel.from_pretrained(original_codeT5, 'Sarthak-506/CodeT5-Fine-Tuned-v1', is_trainable=False)
fine_tuned_codeT5.eval()

# Load StarCoder
StarCoder_tokenizer = AutoTokenizer.from_pretrained(model_id['StarCoder'])
StarCoder = AutoModelForCausalLM.from_pretrained(
    model_id['StarCoder'], device_map="auto"
)
StarCoder.eval()

# Load CodeLlama
CodeLlama_tokenizer = AutoTokenizer.from_pretrained(model_id['CodeLlama'])
CodeLlama = AutoModelForCausalLM.from_pretrained(
    model_id['CodeLlama'], device_map="auto"
)
CodeLlama.eval()


# Generate docstring with CodeT5
def codeT5_generate_docstring(code, model_id):
    if "Base" in model_id:
        model = base_codeT5
    else:
        model = fine_tuned_codeT5

    tokenized_input = codeT5_tokenizer(
        code,
        padding="max_length",
        truncation=True,
        max_length=512,
        return_tensors="pt"
    ).to(device)

    output = model.generate(
        input_ids=tokenized_input['input_ids'],
        attention_mask=tokenized_input['attention_mask'],
        num_beams=5,
        length_penalty=1.0,
        early_stopping=True
    )

    return codeT5_tokenizer.decode(output[0], skip_special_tokens=True)


# StarCoder docstring generation
def StarCoder_generate_docstring(code):
    prompt = f"{code}\n# The goal of this function is to"
    tokenized_input = StarCoder_tokenizer(prompt, return_tensors="pt").to(device)

    output = StarCoder.generate(
        input_ids=tokenized_input['input_ids'],
        attention_mask=tokenized_input['attention_mask'],
        max_new_tokens=40
    )

    return StarCoder_tokenizer.decode(output[0], skip_special_tokens=True)


# CodeLlama docstring generation
def CodeLlama_generate_docstring(code):
    prompt = f"{code}\n# The goal of this function is to"
    tokenized_input = CodeLlama_tokenizer(prompt, return_tensors="pt").to(device)

    output = CodeLlama.generate(
        input_ids=tokenized_input['input_ids'],
        attention_mask=tokenized_input['attention_mask'],
        max_new_tokens=40
    )

    return CodeLlama_tokenizer.decode(output[0], skip_special_tokens=True)


# Selector logic
def provide_docstring(code, model_id):
    if "CodeT5" in model_id:
        return codeT5_generate_docstring(code, model_id)
    elif "StarCoder" in model_id:
        return StarCoder_generate_docstring(code)
    else:
        return CodeLlama_generate_docstring(code)


# Gradio UI
demo = gr.Interface(
    fn=provide_docstring,
    inputs=[
        gr.Textbox(lines=6, label="Enter Code"),
        gr.Dropdown(
            label="Select Model",
            choices=["CodeT5 Base Model", "CodeT5 Fine-Tuned Model", "StarCoder", "CodeLlama"],
            value="CodeT5 Fine-Tuned Model"
        )
    ],
    outputs=gr.Text(label="Generated Docstring"),
    title="Mid-Term Review Demonstration",
    description="Select among different models to generate and compare docstrings for code."
)

demo.launch(share=True)