# Import necessary libraries
from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr
import json
from datetime import datetime

# Load the GPT-2 model and tokenizer from Hugging Face
tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")

# Utility function for generating responses using the GPT-2 model
def generate_response(messages, max_tokens=500, temperature=0.7):
    """
    Generate a response from the model based on the input messages.
    
    Parameters:
    - messages: List of dictionaries containing the role and content of each message.
    - max_tokens: Maximum number of tokens to generate.
    - temperature: Controls randomness in the output.
    
    Returns:
    - The generated response as a string.
    """
    # Concatenate messages into a single prompt
    prompt = "\n".join([f"{msg['role']}: {msg['content']}" for msg in messages])
    
    # Tokenize the input prompt
    input_ids = tokenizer.encode(prompt, return_tensors='pt')
    
    # Generate response
    output = model.generate(input_ids, max_length=len(input_ids[0]) + max_tokens, 
                            temperature=temperature, pad_token_id=tokenizer.eos_token_id)
    
    # Decode the output to a string
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    
    # Return the generated response, excluding the input prompt for clarity
    return response[len(prompt):].strip()

# Function to process user input and generate a response
def process_user_message(user_input, all_messages, debug=True):
    """
    Process the user message and generate a response.
    
    Parameters:
    - user_input: The input from the user.
    - all_messages: A list of previous messages in the conversation.
    - debug: Whether to enable debug logging.
    
    Returns:
    - The response from the model and the updated message history.
    """
    # Add the user's message to the conversation history
    all_messages.append({'role': 'user', 'content': user_input})
    
    # Define a system message for context
    system_message = {
        'role': 'system', 
        'content': "You are a helpful assistant. Answer the user's question as accurately as possible."
    }
    
    # Include the system message and conversation history
    messages = [system_message] + all_messages
    
    # Generate a response using the model
    response = generate_response(messages, max_tokens=500, temperature=0.7)
    
    # Add the model's response to the conversation history
    all_messages.append({'role': 'assistant', 'content': response})
    
    # If debug is enabled, print the conversation history
    if debug:
        print("Conversation History:")
        for msg in all_messages:
            print(f"{msg['role']}: {msg['content']}")
    
    return response, all_messages

# Function to log the messages to a JSON file
def log_messages(new_element, filepath='messages_log.json'):
    try:
        with open(filepath, "r") as file:
            if file.read().strip() == "":
                data = []
            else:
                file.seek(0)
                data = json.load(file)
    except (FileNotFoundError, json.JSONDecodeError):
        data = []
    
    data.append(new_element)
    
    with open(filepath, "w") as file:
        json.dump(data, file, indent=4)

# Initialize an empty list to keep track of the conversation history
context = []

# Function to collect and process user messages
def collect_messages_en(input_text, debug=True):
    global context

    if debug: 
        print(f"User Input: {input_text}")
    if input_text == "":
        return
    
    # Process the user input and generate a response
    response, context = process_user_message(input_text, context, debug=debug)
    
    # Get the current timestamp
    current_timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    # Log the messages
    log_messages({'time_stamp': current_timestamp, 'user_input': input_text, 'AI_response': response})
    
    return response

# Create a Gradio interface for the assistant
demo = gr.Interface(
    fn=collect_messages_en, 
    inputs=gr.Textbox(lines=3, label="Inquiries", placeholder="Ask us anything..."),
    outputs="text",
    title="Customer Service Assistant",
    description="Ask questions about products or services.",
)

demo.launch()