import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

# Define a class named `CodeGenerator` that will be responsible for generating code based on a given prompt.
class CodeGenerator:
    # The constructor initializes the CodeGenerator object with a pre-trained model name.
    # The default model name is "Salesforce/codet5-base".
    def __init__(self, model_name="Salesforce/codet5-base"):
        # Load the pre-trained tokenizer from the specified model name.
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        # Load the pre-trained sequence-to-sequence language model from the specified model name.
        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

    # This method generates code based on the given prompt.
    # The method takes two parameters: `prompt` (the input text) and `max_length` (the maximum length of the generated code).
    def generate_code(self, prompt, max_length=100):
        # Encode the prompt into input IDs that the model can understand.
        input_ids = self.tokenizer.encode(prompt, return_tensors="pt")
        # Generate the output sequence using the pre-trained model.
        # The `generate` method takes the input IDs, the maximum length of the output, and the number of output sequences to return (in this case, 1).
        output = self.model.generate(input_ids, max_length=max_length, num_return_sequences=1)
        # Decode the output sequence and return the generated code.
        return self.tokenizer.decode(output[0], skip_special_tokens=True)

# Define a class named `ChatHandler` that will be responsible for managing the chat history.
class ChatHandler:
    # The constructor initializes the ChatHandler object with an empty chat history.
    def __init__(self):
        self.history = []

    # This method handles incoming messages and generates responses using the provided CodeGenerator.
    # The method takes two parameters: `message` (the user's input message) and `code_generator` (an instance of the CodeGenerator class).
    def handle_message(self, message, code_generator):
        # Generate the response using the provided CodeGenerator.
        response = code_generator.generate_code(message)
        # Append the message-response pair to the chat history.
        self.history.append((message, response))
        # Return the empty message input and the updated chat history.
        return "", self.history

# Define a function named `create_gradio_interface` that creates a Gradio interface for the chat application.
def create_gradio_interface():
    # Create an instance of the CodeGenerator class.
    code_generator = CodeGenerator()
    # Create an instance of the ChatHandler class.
    chat_handler = ChatHandler()

    # Create a Gradio Blocks interface with a soft theme.
    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        # Display a Markdown title for the chat interface.
        gr.Markdown("# S-Dreamer Salesforce/codet5-base Chat Interface")

        # Create a row with two columns.
        with gr.Row():
            # The first column will contain the chat interface.
            with gr.Column(scale=3):
                # Create a chatbot component to display the chat history.
                chatbot = gr.Chatbot(height=400)
                # Create a textbox for the user to input their message.
                message_input = gr.Textbox(label="Enter your code-related query", placeholder="Type your message here...")
                # Create a submit button to send the message.
                submit_button = gr.Button("Submit")

            # The second column will contain the features.
            with gr.Column(scale=1):
                # Display a Markdown title for the features section.
                gr.Markdown("## Features")
                # Define a list of features.
                features = ["Code generation", "Code completion", "Code explanation", "Error correction"]
                # Display each feature as a Markdown list item.
                for feature in features:
                    gr.Markdown(f"- {feature}")
                # Create a button to clear the chat history.
                clear_button = gr.Button("Clear Chat")

        # Connect the submit button to the `handle_message` method of the ChatHandler.
        submit_button.click(chat_handler.handle_message, inputs=[message_input], outputs=[message_input, chatbot])
        # Connect the clear button to a function that clears the chat history.
        clear_button.click(lambda: None, outputs=[chatbot], inputs=[])

    # Launch the Gradio interface.
    demo.launch()

# This is the entry point of the application.
if __name__ == "__main__":
    # Call the `create_gradio_interface` function to start the chat application.
    create_gradio_interface()