from transformers import BertTokenizer, BertForSequenceClassification
import torch
from simple_salesforce import Salesforce
from model import label_encoder  # Assuming the label_encoder was saved in the model.py file

# Load pre-trained model and tokenizer
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
model = BertForSequenceClassification.from_pretrained("results")  # Path to your trained model

# Salesforce authentication setup
sf = Salesforce(username="your_username", password="your_password", security_token="your_token")

# Function to classify the intent of user input
def classify_intent(user_input):
    # Tokenize the user input and make a prediction
    encoded_input = tokenizer(user_input, padding=True, truncation=True, return_tensors="pt")
    
    # Run the model without gradients for efficiency
    with torch.no_grad():
        output = model(**encoded_input)
    
    # Get the predicted label and return it
    predicted_label = torch.argmax(output.logits, dim=1).item()
    return label_encoder.inverse_transform([predicted_label])[0]  # Return intent label

# Function to generate a response based on the classified intent
def generate_response(intent):
    # Simple responses based on predefined intents
    responses = {
        "greeting": "Hi, how can I help you today?",
        "goodbye": "Goodbye! Have a nice day!",
        "question": "Hugging Face is a company that specializes in NLP models."
    }
    
    # Return response for the given intent
    return responses.get(intent, "Sorry, I didn't understand that.")

# Function to create a case in Salesforce based on user input and chatbot response
def create_case_in_salesforce(user_input, response):
    case_data = {
        'Subject': f"User Query: {user_input}",
        'Description': f"User asked: {user_input}\nResponse: {response}"
    }
    
    # Create the case in Salesforce
    case = sf.Case.create(case_data)
    print(f"Case Created with ID: {case['id']}")

# Main function to handle the full interaction
def chatbot_response(user_input):
    # Step 1: Classify the user input into an intent
    intent = classify_intent(user_input)
    
    # Step 2: Generate the appropriate response based on the intent
    response = generate_response(intent)
    
    # Step 3: Optionally create a case in Salesforce (e.g., for queries or support requests)
    create_case_in_salesforce(user_input, response)
    
    # Step 4: Return the chatbot's response
    return response

# Example of using the chatbot
if __name__ == "__main__":
    user_input = "Hello, tell me about Hugging Face."
    print(chatbot_response(user_input))  # Should print the chatbot's response and create a Salesforce case