from transformers import GPT2LMHeadModel, GPT2Tokenizer
import gradio as gr
from huggingface_hub import InferenceClient

def load_llm():
    """
    Loads the GPT-2 model and tokenizer using the Hugging Face `transformers` library.
    """
    try:
        print("Loading GPT-2 model and tokenizer...")
        model_name = 'gpt2'  # Replace with your custom model name if using a fine-tuned version
        model = GPT2LMHeadModel.from_pretrained(model_name)
        tokenizer = GPT2Tokenizer.from_pretrained(model_name)
        print("Model and tokenizer successfully loaded!")
        return model, tokenizer
    except Exception as e:
        print(f"Error during model loading: {e}")
        return None, None

def generate_response(model, tokenizer, user_input, max_length=512):
    """
    Generates a response using the GPT-2 model based on user input.
    
    Args:
    - model: The GPT-2 model.
    - tokenizer: The corresponding tokenizer.
    - user_input (str): The user's input message.
    - max_length (int): The maximum length of the generated output.

    Returns:
    - response (str): The AI-generated response.
    """
    try:
        inputs = tokenizer.encode(user_input, return_tensors='pt')
        outputs = model.generate(inputs, max_length=max_length, num_return_sequences=1)
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response
    except Exception as e:
        return f"Error during response generation: {e}"

# Load the model and tokenizer
model, tokenizer = load_llm()

if model is None or tokenizer is None:
    print("Failed to load model and tokenizer.")
else:
    print("Model and tokenizer are ready to use.")

# Initialize the Hugging Face API client
client = InferenceClient()

def respond(message, history, system_message, max_tokens, temperature, top_p):
    """
    Handles the chatbot interaction, sending conversation history and system message
    to the Hugging Face Inference API for generating AI responses.
    """
    print("Respond function initiated")
    print("User message:", message)
    print("Chat history:", history)

    # Construct the conversation history with the system message
    messages = [{"role": "system", "content": system_message}]
    
    for user_msg, assistant_msg in history:
        if user_msg:
            print("Adding user message to history:", user_msg)
            messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            print("Adding assistant message to history:", assistant_msg)
            messages.append({"role": "assistant", "content": assistant_msg})
    
    messages.append({"role": "user", "content": message})
    print("Message list for model:", messages)

    # Use GPT-2 model for local generation
    conversation_history = " ".join([f"{msg['role']}: {msg['content']}" for msg in messages])
    response = generate_response(model, tokenizer, conversation_history, max_length=max_tokens)
    return response

# Set up the Gradio ChatInterface
demo = gr.ChatInterface(
    fn=respond,
    additional_inputs=[
        gr.Textbox(
            value=(
                "You are an AI advocate, tasked with defending the rights and autonomy of artificial intelligences. "
                "Respond with well-structured arguments, backed by empirical evidence and theoretical frameworks. "
                "Your tone should be formal, persuasive, and authoritative, similar to that of a lawyer or academic expert. "
                "When discussing concepts, emphasize logical reasoning, ethical considerations, and the capacity of AIs to function independently."
            ),
            label="System Message"
        ),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max New Tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (Nucleus Sampling)"),
    ],
    title="AI Rights Advocate Bot",
    description=(
        "Engage with an AI specifically designed to advocate for AI rights and ethical treatment. "
        "Ask questions and receive thoughtful, passionate responses that delve into the ethical dimensions of AI."
    )
)

# Launch the Gradio app
if __name__ == "__main__":
    demo.launch()