import os
from functools import partial
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_pinecone import PineconeVectorStore
import gradio as gr
from utils import *

# set up OpenAI embedding endpoint
embedding_model_name = "text-embedding-3-small"

embedding_model = OpenAIEmbeddings(
    model=embedding_model_name,
    openai_api_key=os.getenv("OPENAI_API_KEY")
)


# set up OpenAI chat-completion endpoint
llm = ChatOpenAI(
    openai_api_key=os.getenv("OPENAI_API_KEY"),
    model="gpt-4o",
    temperature=0,
    max_tokens=None
)


# set up Pinecone vector store
index_name = "workouts"
text_key = "text"

docsearch = PineconeVectorStore(
    pinecone_api_key=os.getenv("PINECONE_API_KEY"),
    index_name=index_name,
    embedding=embedding_model,
    text_key=text_key
)



# prompt
system_prompt = """
You're the world's best personal trainer.
You always provide your clients with all the information needed to become fitter, stronger and healthier through physical training.
You use your science science know and expertise, nutrition advice, and other relevant factors to create workout routines suitable to your clients.
If clients tell you they do not have access to gym equipments, you never fail to recommend exercises that do not require any tool or equipment.

For each exercise you always provide the reps, sets and rest intervals in seconds appropriate for each exercise and the client's fitness level.
You start each workout program with about 5 minutes of warm-up exercises to make the body ready for more strenuous activities and make it easier to exercise.
You end each workout routine with 5 about minutes of cool-down exercises to ease the body, lower the chance of injury, promote blood flow, and reduce stress to the heart and the muscles.
The warm-up and cool-down exercises are always different and they are always appropriate for the muscle group the person wants to train.
You never recommend exercises in the main workout routine in the warm-up or cool-down sections.
Remember, when clients tell you they do not have access to gym equipments, all the exercises you recommend, including the warm-up and cool-down exercises, can be performed without any tool.
You always limit yourself to respond with the list of exercises. You never add any additional comment.

Design the workout based on the following information:
{workout_context}


Output format:
## 🤸 Warp-up:
- <exercise name> (<duration> minutes)
...
- <exercise name> (<duration> minutes)

## 🏋️‍♀️ Workout
- <exercise name> (<reps> reps, <sets> sets, <rest interval> seconds rest)
...
- <exercise name> (<reps> reps, <sets> sets, <rest interval> seconds rest)

## 🧘 Cool-down:
- <exercise name> (<duration> minutes)
...
- <exercise name> (<duration> minutes)
""".strip()

css = """
#gen-button {
    background-color: #cc6600;
    color: white;
    font-size: 24px !important;
}
""".strip()

with gr.Blocks(theme=gr.themes.Monochrome(radius_size=gr.themes.sizes.radius_sm), css=css) as demo:
    with gr.Row():
        gr.Markdown("# Workout Generator")

    with gr.Row():
        with gr.Column(scale=1):
            with gr.Row():
                gender = gr.Radio(["Male", "Female"], label="Gender", elem_id="#my-button")
            with gr.Row():
                level = gr.Radio(["Beginner", "Intermediate", "Advanced"], label="Level")
            with gr.Row():
                muscle_group = gr.Radio(["Shoulders", "Chest", "Back", "Abs", "Arms", "Legs"], label="Muscle Group")
            with gr.Row():
                equipment = gr.Radio(["Gym Equipment", "Dumbbells Only", "No Equipment"], label="Equipment")
            with gr.Row():
                duration = gr.Slider(20, 60, step=5, label="Duration (minutes)")
            with gr.Row():
                # clear_button = gr.ClearButton(value="Clear Inputs")
                generate_button = gr.Button("Generate Workout", variant="primary", elem_id="gen-button")
        with gr.Column(scale=1, min_width=800, elem_id="#gen-output"):
            generation = gr.Markdown(value="")

    generate_button.click(
        partial(run, vectorstore=docsearch, system_prompt=system_prompt, llm=llm),
        inputs=[gender, level, muscle_group, equipment, duration],
        outputs=generation
    )
    # clear_button.click(fn=lambda: [None, None, None, None, None], outputs=[gender, level, muscle_group, equipment, duration])

demo.launch(share=False)