import gradio as gr
from transformers import pipeline

# Define your Minecraft actions with descriptions
ACTIONS = {
    "move_towards_player": "walk or move toward a specific player",
    "stand_still": "remain in place, don't move",
    "follow_player": "continuously follow behind a player",
    "run_away": "move away from a player or danger",
    "jump": "jump up or over something",
    "mine_block": "break or mine a block",
    "place_block": "place or put down a block",
    "attack": "attack or fight something",
    "use_item": "use or activate an item",
    "chat_only": "just chat without taking physical action"
}

# Load the zero-shot classification pipeline
# Using a model that works well with Spaces and has good zero-shot performance
classifier = pipeline("zero-shot-classification", 
                      model="facebook/bart-large-mnli")

def predict_action(text):
    # Extract just the action names for classification
    action_names = list(ACTIONS.keys())
    
    # Add context to improve classification by combining action names with descriptions
    candidate_labels = [f"{action}: {ACTIONS[action]}" for action in action_names]
    
    # Run zero-shot classification
    result = classifier(text, candidate_labels, multi_label=False)
    
    # Extract results
    action_scores = {}
    for i, label in enumerate(result["labels"]):
        # Extract just the action name from the combined label+description
        action_name = label.split(":")[0]
        action_scores[action_name] = result["scores"][i]
    
    # Find the best action
    best_action = max(action_scores, key=action_scores.get)
    confidence = action_scores[best_action]
    
    return {
        "action": best_action,
        "confidence": confidence,
        "all_actions": action_scores
    }

# Example to test (you can remove this in production)
example_text = "I want to go and meet that player over there"
test_result = predict_action(example_text)
print(f"Example: '{example_text}'")
print(f"Best action: {test_result['action']}")
print(f"Confidence: {test_result['confidence']:.4f}")

# Create the Gradio interface
demo = gr.Interface(
    fn=predict_action,
    inputs=gr.Textbox(label="Character text from C.AI"),
    outputs=gr.JSON(label="Predicted Action"),
    title="Minecraft NPC Action Predictor",
    description="""This tool analyzes what your C.AI character says and determines the best action to take in Minecraft.
                   Input the character's dialogue to get the appropriate in-game action.""",
    examples=[
        ["I'll help you build that house!"],
        ["Look out, there's a creeper behind you!"],
        ["Let me show you the way to the village."],
        ["I think I'll just wait here for a while."],
        ["I need to break this stone to make a path."]
    ]
)

demo.launch()