"""
Workout Session Data Collection
Gradio app for collecting real-world labeled workout data.
Submissions are written to a private Google Sheet in real time.
"""

import gradio as gr
import gspread
import os
import json
from datetime import datetime
from google.oauth2.service_account import Credentials

# ─────────────────────────────────────────────────────────────
# GOOGLE SHEETS SETUP
# Credentials are loaded from the GOOGLE_CREDENTIALS env var
# (set as a HF Space Secret — paste the entire service account
# JSON as a single-line string)
# ─────────────────────────────────────────────────────────────

SCOPES = [
    "https://www.googleapis.com/auth/spreadsheets",
    "https://www.googleapis.com/auth/drive",
]

SHEET_NAME = os.getenv("SHEET_NAME", "workout_sessions")

# Expected header row — must match the sheet's first row exactly
HEADERS = [
    "timestamp",
    "user_text",
    "mood",
    "exertion",
    "soreness_region",
    "soreness_severity",
    "completion_status",
]


def get_sheet():
    """
    Authenticate with Google Sheets using the service account
    credentials stored in the GOOGLE_CREDENTIALS env var.
    Returns the first worksheet of the target spreadsheet.
    """
    creds_json = os.getenv("GOOGLE_CREDENTIALS")
    if not creds_json:
        raise EnvironmentError(
            "GOOGLE_CREDENTIALS env var not set. "
            "Add your service account JSON as a HF Space Secret."
        )

    creds_dict = json.loads(creds_json)
    creds      = Credentials.from_service_account_info(creds_dict, scopes=SCOPES)
    client     = gspread.authorize(creds)
    sheet      = client.open(SHEET_NAME).sheet1

    # Write headers if the sheet is empty
    if sheet.row_count == 0 or sheet.row_values(1) != HEADERS:
        sheet.insert_row(HEADERS, index=1)

    return sheet


# ─────────────────────────────────────────────────────────────
# LABEL OPTIONS
# ─────────────────────────────────────────────────────────────

MOOD_OPTIONS = [
    "accomplished", "anxious", "distracted", "energized",
    "fatigued",     "frustrated", "neutral",  "positive",
]

EXERTION_OPTIONS = ["low", "moderate", "high"]

SORENESS_REGION_OPTIONS = [
    "none", "back", "biceps", "chest",
    "legs", "shoulder", "triceps",
]

SORENESS_SEVERITY_OPTIONS = ["none", "mild", "moderate", "severe"]

COMPLETION_OPTIONS = ["full", "partial"]

# ─────────────────────────────────────────────────────────────
# NUMERICAL ENCODING MAPS
# Must stay in sync with your training notebook label maps
# ─────────────────────────────────────────────────────────────

MOOD_MAP = {
    "accomplished": 0,
    "anxious":      1,
    "distracted":   2,
    "energized":    3,
    "fatigued":     4,
    "frustrated":   5,
    "neutral":      6,
    "positive":     7,
}

EXERTION_MAP = {
    "low":      0,
    "moderate": 1,
    "high":     2,
}

SORENESS_REGION_MAP = {
    "none":     0,
    "back":     1,
    "biceps":   2,
    "chest":    3,
    "legs":     4,
    "shoulder": 5,
    "triceps":  6,
}

SORENESS_SEVERITY_MAP = {
    "none":     0,
    "mild":     1,
    "moderate": 2,
    "severe":   3,
}

COMPLETION_MAP = {
    "full":    1,
    "partial": 0,
}


# ─────────────────────────────────────────────────────────────
# SUBMISSION HANDLER
# ─────────────────────────────────────────────────────────────

def submit_entry(
    user_text:         str,
    mood:              str,
    exertion:          str,
    soreness_region:   str,
    soreness_severity: str,
    completion_status: str,
) -> str:
    """
    Validates the form and appends one row to the Google Sheet.
    Returns a status string displayed to the user.
    """

    # ── Validation ───────────────────────────────────────────
    if not user_text or len(user_text.strip()) < 10:
        return "⚠️ Please describe your session in at least 10 characters."

    missing = []
    if not mood:              missing.append("Mood")
    if not exertion:          missing.append("Exertion")
    if not soreness_region:   missing.append("Soreness Region")
    if not soreness_severity: missing.append("Soreness Severity")
    if not completion_status: missing.append("Completion")

    if missing:
        return f"⚠️ Please select: {', '.join(missing)}"

    # Severity/region consistency check
    if soreness_region == "none" and soreness_severity != "none":
        return "⚠️ Soreness severity must be 'none' when region is 'none'."
    if soreness_region != "none" and soreness_severity == "none":
        return "⚠️ Please select a severity level for your soreness region."

    # ── Encode labels to integers ─────────────────────────────
    row = [
        datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S"),
        user_text.strip(),
        MOOD_MAP[mood],
        EXERTION_MAP[exertion],
        SORENESS_REGION_MAP[soreness_region],
        SORENESS_SEVERITY_MAP[soreness_severity],
        COMPLETION_MAP[completion_status],
    ]

    # ── Write to Google Sheets ───────────────────────────────
    try:
        sheet = get_sheet()
        sheet.append_row(row, value_input_option="USER_ENTERED")
    except EnvironmentError as e:
        return f"❌ Configuration error: {str(e)}"
    except Exception as e:
        return f"❌ Failed to save: {str(e)}"

    return (
        "✅ Submitted! Thank you — your session has been logged.\n\n"
        "Feel free to submit another entry."
    )


# ─────────────────────────────────────────────────────────────
# GRADIO UI
# ─────────────────────────────────────────────────────────────

with gr.Blocks(
    title="Workout Session Logger",
    theme=gr.themes.Soft(
        primary_hue="orange",
        neutral_hue="slate",
    ),
    css="""
        .submit-btn { background: #FF4500 !important; border: none !important; }
        .submit-btn:hover { background: #CC3700 !important; }
        footer { display: none !important; }
    """,
) as demo:

    # ── Header ────────────────────────────────────────────────
    gr.Markdown(
        """
        # 🏋️ Workout Session Logger
        ### Help train a smarter fitness AI
        Describe your session in your own words, then label it using the
        dropdowns below. Your data helps build a model that understands
        how athletes really feel after training.

        **All submissions are anonymous. Be as honest as possible —
        bad sessions are just as valuable as great ones.**
        """
    )

    gr.Markdown("---")

    # ── Free text input ───────────────────────────────────────
    gr.Markdown("### 1. Describe your session")
    gr.Markdown(
        "_Write naturally — exactly how you'd tell a training partner. "
        "Include how you felt, what you trained, any soreness or energy levels._"
    )

    user_text = gr.Textbox(
        label="Session description",
        placeholder=(
            "e.g. Hit a new PR on deadlifts today, feeling absolutely "
            "wrecked but stoked. Lower back is pretty sore and I only "
            "got through 4 of 5 sets before calling it..."
        ),
        lines=5,
        max_lines=10,
    )

    gr.Markdown("---")

    # ── Labels ────────────────────────────────────────────────
    gr.Markdown("### 2. Label your session")
    gr.Markdown(
        "_Select the option that best matches how you felt **after** the session._"
    )

    with gr.Row():
        mood = gr.Dropdown(
            choices=MOOD_OPTIONS,
            label="Mood",
            info="How did you feel after finishing?",
        )
        exertion = gr.Dropdown(
            choices=EXERTION_OPTIONS,
            label="Exertion level",
            info="How hard did you push overall?",
        )

    with gr.Row():
        soreness_region = gr.Dropdown(
            choices=SORENESS_REGION_OPTIONS,
            label="Soreness region",
            info="Which muscle group is most sore? Select 'none' if no soreness.",
        )
        soreness_severity = gr.Dropdown(
            choices=SORENESS_SEVERITY_OPTIONS,
            label="Soreness severity",
            info="How intense is the soreness? Select 'none' if no soreness.",
        )

    completion_status = gr.Dropdown(
        choices=COMPLETION_OPTIONS,
        label="Completion",
        info="Did you finish the full planned session?",
    )

    gr.Markdown("---")

    # ── Submit ────────────────────────────────────────────────
    submit_btn = gr.Button(
        "Submit Session",
        variant="primary",
        elem_classes=["submit-btn"],
        size="lg",
    )

    status_output = gr.Textbox(
        label="Status",
        interactive=False,
        show_label=False,
    )

    submit_btn.click(
        fn=submit_entry,
        inputs=[
            user_text,
            mood,
            exertion,
            soreness_region,
            soreness_severity,
            completion_status,
        ],
        outputs=status_output,
    )

    # ── Label reference ───────────────────────────────────────
    with gr.Accordion("📖 Label reference guide", open=False):
        gr.Markdown(
            """
            ### Mood
            | Label | When to use |
            |---|---|
            | **accomplished** | Hit a goal, PR, or felt proud of the effort |
            | **energized** | Left the gym feeling charged and strong |
            | **positive** | Good session, happy with the work done |
            | **neutral** | Average session, nothing special either way |
            | **fatigued** | Drained, low energy, struggled through it |
            | **frustrated** | Bad session, missed lifts, things went wrong |
            | **anxious** | Felt uneasy, worried about injury or performance |
            | **distracted** | Couldn't focus, mind elsewhere, scattered session |

            ### Exertion
            | Label | When to use |
            |---|---|
            | **low** | Light session, easy pace, well within limits |
            | **moderate** | Solid effort, challenging but manageable |
            | **high** | Pushed to near-limit, very demanding session |

            ### Soreness Region
            Select the **primary** muscle group that is sore or was most targeted.
            Choose **none** if you have no notable soreness.

            ### Soreness Severity
            | Label | When to use |
            |---|---|
            | **none** | No soreness at all |
            | **mild** | Slightly tight or tender, barely noticeable |
            | **moderate** | Noticeably sore, aware of it during movement |
            | **severe** | Very sore, impacts range of motion or daily activity |

            ### Completion
            | Label | When to use |
            |---|---|
            | **full** | Completed all planned sets, exercises, and volume |
            | **partial** | Skipped exercises, cut sets short, or left early |
            """
        )


# ─────────────────────────────────────────────────────────────
# LAUNCH
# ─────────────────────────────────────────────────────────────

if __name__ == "__main__":
    demo.launch()