"""
app.py
------
Streamlit frontend for the coaching recommendation system.
"""

import os
import streamlit as st
from setup_db import setup
from retrieval import retrieve
from reranker import rerank

# ── Page config ───────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="Coaching Assistant",
    page_icon="💪🏻",
    layout="centered"
)

# ── One-time DB setup ─────────────────────────────────────────────────────────
# Sets up the SQLite DB if it doesn't already exist.
if not os.path.exists("exercises.db"):
    with st.spinner("Setting up exercise database..."):
        setup()

# ── Header ────────────────────────────────────────────────────────────────────
st.title("💪🏻 Coaching Assistant")
st.markdown(
    "Describe what you need and get personalised exercise recommendations. "
    "You can mention goals, injuries, equipment, or intensity."
)
st.divider()

# ── Example queries ───────────────────────────────────────────────────────────
st.markdown("**Try an example:**")
examples = [
    "I have knee pain and need low-impact exercises",
    "Explosive drills for a winger",
    "Upper body rehab, no weights",
    "Beginner core strengthening",
]
cols = st.columns(len(examples))
for i, example in enumerate(examples):
    if cols[i].button(example, use_container_width=True):
        st.session_state["query"] = example

# ── Query input ───────────────────────────────────────────────────────────────
query = st.text_input(
    label="Your query",
    placeholder="e.g. low-impact knee rehab for a beginner",
    key="query"
)

search_clicked = st.button("Get Recommendations", type="primary", use_container_width=True)

# ── Pipeline ──────────────────────────────────────────────────────────────────
if search_clicked and query.strip():

    with st.spinner("Finding the best exercises for you..."):

        # Stage 1: BM25 retrieval
        candidates = retrieve(query, top_k=15)

        if not candidates:
            st.warning("No matching exercises found. Try different keywords.")
            st.stop()

        # Stage 2: LLM re-ranking
        try:
            results = rerank(query, candidates, top_n=5)
        except Exception as e:
            st.error(f"Re-ranking failed: {e}")
            # Graceful fallback: show top BM25 results without re-ranking
            results = [{**c, "reason": ""} for c in candidates[:5]]

    # ── Results ───────────────────────────────────────────────────────────────
    st.divider()
    st.subheader(f"Top recommendations for: *{query}*")
    st.caption(f"Retrieved {len(candidates)} candidates → re-ranked to top {len(results)}")

    for i, ex in enumerate(results, 1):

        # Difficulty badge colour
        diff_color = {"beginner": "🟢", "intermediate": "🟡", "advanced": "🔴"}.get(
            ex.get("difficulty", "").lower(), "⚪"
        )

        with st.container(border=True):
            col1, col2 = st.columns([3, 1])
            with col1:
                st.markdown(f"**{i}. {ex['title']}**")
                st.caption(ex.get("description", ""))
            with col2:
                st.markdown(f"{diff_color} {ex.get('difficulty', '').capitalize()}")
                st.caption(f"🏷 {ex.get('body_part', '').capitalize()}")

            if ex.get("reason"):
                st.info(f"💡 {ex['reason']}")

            # Metadata pills
            meta_cols = st.columns(3)
            meta_cols[0].markdown(f"**Equipment:** {ex.get('equipment', 'none') or 'none'}")
            meta_cols[1].markdown(f"**Intensity:** {ex.get('intensity', '—')}")
            meta_cols[2].markdown(f"**Focus:** {ex.get('injury_focus', '—') or '—'}")

elif search_clicked and not query.strip():
    st.warning("Please enter a query first.")

# ── Footer ────────────────────────────────────────────────────────────────────
st.divider()
st.caption(
    "Pipeline: BM25 keyword retrieval → Mistral-7B re-ranking | "
    "Built with Streamlit · Deployed on HuggingFace Spaces · 100% free"
)