import gradio as gr
import difflib
import os

# A simple mapping of prompt keywords to image file names
exercise_sets = {
    "shoulder": "shoulder_mobility.png",
    "neck": "neck_stretches.png",
    "back": "back_mobility.png",
    "full body": "full_body_mobility.png",
    "warm up": "warmup_routine.png"
}

# Fuzzy matching prompt to known categories
def match_prompt_to_exercise(prompt):
    keywords = list(exercise_sets.keys())
    closest = difflib.get_close_matches(prompt.lower(), keywords, n=1, cutoff=0.3)
    if closest:
        return exercise_sets[closest[0]]
    else:
        return None

# Main function
def get_exercise_set(prompt):
    matched_file = match_prompt_to_exercise(prompt)
    if matched_file and os.path.exists(f"exercise_sets/{matched_file}"):
        return f"exercise_sets/{matched_file}", f"exercise_sets/{matched_file}"
    else:
        return None, "Sorry, no matching set found. Try: 'shoulder', 'neck', 'warm up'"

# Interface
gr.Interface(
    fn=get_exercise_set,
    inputs=gr.Textbox(label="Enter your goal (e.g., 'Shoulder mobility exercises')"),
    outputs=[gr.Image(label="Generated Exercise Sheet"), gr.File(label="Download Image")],
    title="🧘 AI Exercise Sheet Generator",
    description="Type a goal like 'Neck stretches' or 'Warm-up routine' to get an illustrated exercise set.",
    examples=[
        ["Shoulder mobility exercises"],
        ["Neck stretches"],
        ["Warm-up before boxing"],
        ["Full body mobility"]
    ]
).launch()