import streamlit as st
import pandas as pd
from auth import generate_otp, create_jwt, verify_jwt
from email_utils import send_otp_email
from dotenv import load_dotenv
from database import add_user, user_exists, verify_login

# Load environment variables
load_dotenv()

# --- Page Config ---
st.set_page_config(
    page_title="FitPlan AI",
    page_icon="🏋️‍♂️",
    layout="wide",
    initial_sidebar_state="expanded"
)

# --- INNOVATIVE LIGHT MODE UI ---
st.markdown("""
<style>
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@200;300;400;500;600;700;800;900&family=Orbitron:wght@400;700;900&display=swap');

    :root {
        --bg-primary: linear-gradient(135deg, #f8fafc 0%, #f1f5f9 50%, #e2e8f0 100%);
        --glass-bg: rgba(255, 255, 255, 0.85);
        --glass-border: rgba(148, 163, 184, 0.2);
        --primary-gradient: linear-gradient(135deg, #6366f1 0%, #8b5cf6 50%, #a855f7 100%);
        --success-gradient: linear-gradient(135deg, #10b981 0%, #34d399 100%);
        --text-primary: #0f172a;
        --text-secondary: #475569;
    }

    * {
        font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
    }

    .stApp {
        background: var(--bg-primary);
        color: var(--text-primary);
    }

    .hero-section {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 30%, #f093fb 60%, #f5576c 100%);
        padding: 4rem 3rem;
        border-radius: 32px;
        text-align: center;
        margin-bottom: 4rem;
        box-shadow: 0 0 30px rgba(99, 102, 241, 0.4);
        position: relative;
        overflow: hidden;
    }

    .hero-title {
        font-family: 'Orbitron', monospace;
        font-size: clamp(2.5rem, 6vw, 4.5rem);
        font-weight: 900;
        background: linear-gradient(45deg, #ffffff 0%, #f0f9ff 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        margin: 0 0 1rem 0;
    }

    .hero-subtitle {
        font-size: 1.4rem;
        color: rgba(255,255,255,0.95);
        font-weight: 300;
        margin: 0;
    }

    .glass-card {
        background: var(--glass-bg);
        backdrop-filter: blur(25px) saturate(180%);
        border: 1px solid var(--glass-border);
        border-radius: 24px;
        padding: 2.5rem;
        box-shadow: 0 25px 50px rgba(0,0,0,0.1);
        margin-bottom: 2rem;
        position: relative;
    }

    .glass-card::before {
        content: '';
        position: absolute;
        top: 0;
        left: 0;
        right: 0;
        height: 4px;
        background: var(--primary-gradient);
        border-radius: 24px 24px 0 0;
    }

    .stButton > button {
        background: var(--primary-gradient) !important;
        border: none !important;
        border-radius: 20px !important;
        height: 4rem !important;
        font-weight: 600 !important;
        font-size: 1.1rem !important;
        color: white !important;
        box-shadow: 0 0 30px rgba(99, 102, 241, 0.4) !important;
    }

    .plan-container {
        background: rgba(255, 255, 255, 0.95);
        backdrop-filter: blur(30px);
        padding: 3rem;
        border-radius: 28px;
        border: 2px solid rgba(99, 102, 241, 0.2);
        white-space: pre-wrap;
        font-size: 1.1rem;
        line-height: 1.8;
        box-shadow: 0 30px 60px rgba(0,0,0,0.12);
        position: relative;
        max-height: 70vh;
        overflow-y: auto;
    }

    section[data-testid="stSidebar"] {
    background: linear-gradient(180deg, rgba(255,255,255,0.95) 0%, rgba(248, 250, 252, 0.95) 100%);
    backdrop-filter: blur(20px);
}
</style>
""", unsafe_allow_html=True)

# --- Initialize Session State ---
if "user_data" not in st.session_state:
    st.session_state.user_data = None
if "otp" not in st.session_state:
    st.session_state.otp = None
if "authenticated" not in st.session_state:
    st.session_state.authenticated = False
if "token" not in st.session_state:
    st.session_state.token = None


# --- FULL WORKOUT GENERATOR WITH ALL PLANS ---
def workout_generator(prompt, goal, max_new_tokens=1000):
    plans = {
        "Build Muscle": """
➡️Day 1: Chest (Push)
- Focus on the "stretch and squeeze" of the pectoral fibers.
- Barbell Bench Press: 4 sets 6–8 reps.
- Incline Dumbbell Press: 3 sets 10–12 reps.
- Chest Dips (Leaning forward): 3 sets Failure.
- Cable Flys (Middle to Low): 3 sets 15 reps.
- Push-ups: 2 sets Max reps (as a finisher).

➡️Day 2: Back (Pull)
- Drive with your elbows to ensure the lats are doing the work, not just your biceps.
- Deadlifts: 3 sets 5 reps (Heavy).
- Pull-Ups (or Lat Pulldowns): 4 sets 8–10 reps.
- Bent-Over Barbell Rows: 3 sets 10 reps.
- Seated Cable Rows: 3 sets 12 reps.
- Face Pulls: 3 sets 15–20 reps (for rear delts and posture).

➡️Day 3: Legs (Lower)
- The most demanding day. Ensure you hit full depth on your movements.
- Barbell Back Squats: 4 sets 8–10 reps.
- Leg Press: 3 sets 12–15 reps.
- Romanian Deadlifts: 3 sets 10–12 reps (for hamstrings).
- Leg Extensions: 3 sets 15 reps.
- Seated Calf Raises: 4 sets 15–20 reps.

➡️Day 4: Shoulders
- Prioritize the side (lateral) head of the delt for maximum width.
- Overhead Press (Barbell or DB): 4 sets 8–10 reps.
- Lateral Raises (Dumbbell): 4 sets 15–20 reps.
- Front Raises (Plate or DB): 3 sets 12 reps.
- Reverse Pec Deck (Rear Delts): 3 sets 15 reps.
- Dumbbell Shrugs: 3 sets 12 reps.

➡️Day 5: Arms & Abs
- High volume to force blood into the smaller muscle groups.
- Barbell Curls: 3 sets 10 reps.
- Close-Grip Bench Press: 3 sets 8–10 reps.
- Hammer Curls: 3 sets 12 reps.
- Tricep Rope Pushdowns: 3 sets 15 reps.
- Hanging Leg Raises: 3 sets 15 reps.
- Plank: 3 sets 60 seconds
""",

        "Weight Loss": """
➡️Day 1: Upper Body (Push)
- Focuses on chest, shoulders, and triceps.
- Dumbbell Chest Press: 3 sets of 10–12 reps.
- Shoulder Press: 3 sets of 10–12 reps.
- Triceps Pushdowns: 3 sets of 12–15 reps.
- Push-ups: 3 sets to failure.
- Cardio Finisher: 15 minutes of brisk walking.

➡️Day 2: Lower Body
- Focuses on quads, glutes, and hamstrings.
- Goblet Squats: 3 sets of 12 reps.
- Walking Lunges: 3 sets of 10 reps per leg.
- Romanian Deadlifts: 3 sets of 10 reps.
- Glute Bridges: 3 sets of 15 reps.
- Cardio Finisher: 15 minutes on a cycle or treadmill.

➡️Day 3: Cardio HIIT
- Designed to spike your heart rate and burn fat quickly.
- Format: 30 seconds of high intensity (sprint/jumping jacks/burpees) followed by 60 seconds of low intensity (walking).
- Repeat: 10–15 rounds.
- Exercises to Mix: Mountain climbers, burpees, high knees, and skipping rope.

➡️Day 4: Upper Body (Pull)
- Focuses on the back and biceps.
- Lat Pulldowns: 3 sets of 10–12 reps.
- Seated Cable Rows: 3 sets of 12 reps.
- Dumbbell Bicep Curls: 3 sets of 12 reps.
- Face Pulls: 3 sets of 15 reps.
- Cardio Finisher: 15 minutes on the rowing machine.

➡️Day 5: Steady-State Cardio & Core
- Lower intensity to promote recovery while still burning calories.
- Steady Cardio: 35–40 minutes of jogging, swimming, or brisk walking.
- Core Circuit (Repeat 3 times).
- Plank: Hold for 45–60 seconds.
- Russian Twists: 20 reps.
- Bicycle Crunches: 15 reps per side.
""",

        "Strength Gain": """
➡️Day 1: Full Body Strength
- Focus on hitting every major muscle group in a single, high-intensity session.
- Barbell Squats: 3–5 sets of 5 reps.
- Barbell Bench Press: 3–5 sets of 5 reps.
- Barbell Rows: 3 sets of 6–8 reps.
- Overhead Press: 3 sets of 5–8 reps.
- Deadlifts: 1–3 sets of 5 reps.

➡️Day 2: Lower Body
- Focus on the largest muscles in the body to build massive foundational strength.
- Barbell Squats: 3 sets of 5–8 reps (Main power move).
- Romanian Deadlifts: 3 sets of 8–10 reps (Targets hamstrings and glutes).
- Leg Press or Lunges: 3 sets of 10 reps (Builds quad volume).
- Calf Raises: 3 sets of 15 reps (Lower leg stability).
- Plank / Core Work: 3 sets of 60 seconds.

➡️Day 3: Leg Day
- Hit the main lifts for the second time this week.
- Barbell Squats: 5 sets of 5 reps.
- Barbell Bench Press: 5 sets of 5 reps.
- Barbell Rows: 5 sets of 5 reps.
- Dips: 3 sets to failure (for chest and tricep endurance).
- Hanging Leg Raises: 3 sets of 15 reps (core strength).

➡️Day 4: Rest & Recovery
- Light Walk: 20 mins to boost blood flow.
- Stretch: Focus on hips and lower back.
- Hydrate: Drink extra water for muscle repair.
- Protein: Keep intake high to rebuild tissue.
- Sleep: Aim for 8 hours of quality rest.

➡️Day 5: Upper Body
- Dumbbell Shoulder Press: 3 sets × 12 reps (Focus on vertical control).
- Cable Rows: 3 sets × 12–15 reps (Keep chest up, squeeze shoulder blades).
- Face Pulls: 3 sets × 15–20 reps (Great for posture and rear delts).
- Hammer Curls: 3 sets × 12 reps (Targets forearms and biceps).
- Tricep Overhead Extension: 3 sets × 12 reps
""",

        "Abs Building": """
➡️Day 1: Upper Abs Focus
- Floor Crunch: 3 sets x 20 reps. Targets the upper rectus abdominis.
- Sit-ups: 4 sets x 15 reps. A classic for building core power.
- Toe Touches: 4 sets x 30 reps. Engages the upper abs through targeted flexion.
- Plank: 3 sets x Max Time. Essential for overall core tension.

➡️Day 2: Lower Abs Focus
- Lying Leg Raise: 3 sets x 15 reps. One of the best for lower ab activation.
- Reverse Crunch: 4 sets x 10-15 reps. Focuses on bringing the hips toward the chest.
- Flutter Kicks: 40 seconds. Keeps constant tension on the lower core.
- Scissor Kicks: 4 sets x 30 reps. Mimics walking to engage deep lower muscles.

➡️Day 3: Obliques
- Russian Twists: 4 sets x 20 reps. Key for rotational strength and side-ab definition.
- Side Plank: 30-45 seconds per side. Builds lateral stability.
- Crossbody Mountain Climbers: 30 seconds. Combines cardio with oblique rotation.
- Side Crunches: 40 seconds per side. Directly targets the love handle area.

➡️Day 4: Core Stability
- Bird Dog: 5 reps per side (slowly). Teaches the core to stay stable during movement.
- Plank Shoulder Taps: 40 seconds. Forces the core to resist rotation.
- Superman Hold: 40 seconds. Important for lower back health to support front abs.
- Plank Jacks: 10-15 reps. Adds a cardio element to traditional planking.

➡️Day 5: Total Core Burn
- Bicycle Crunches: 40 seconds. Often cited as the most effective ab exercise.
- V-Ups: 10 reps. An advanced move that hits upper and lower abs simultaneously.
- Mountain Climbers: 30 seconds. Explosive movement for final fat burning.
- High Knees: 30 seconds. Keeps heart rate up to maximize the "afterburn" effect.
""",

        "Flexibility": """
➡️Day 1: Full Body Dynamic Mobility
- Neck Rolls: 2 sets of 30 seconds.
- Arm Circles: 2 sets of 20 reps (forward and backward).
- Cat-Cow Stretch: 3 sets of 10 reps to warm up the spine.
- Leg Swings: 2 sets of 15 reps per leg.
- Standing Forward Fold: 3 sets of 20 seconds.

➡️Day 2: Hips & Hamstrings
- Downward Dog: 3 sets of 30 seconds for calves and hamstrings.
- Lunge Stretch: 3 sets of 20 seconds per leg to open hip flexors.
- Butterfly Pose: 3 sets of 30 seconds for inner thighs.
- Seated Hamstring Stretch: 3 sets of 30 seconds (lean from hips, not waist).
- Pigeon Pose: 2 sets of 30 seconds per side for deep hip opening.

➡️Day 3: Spine & Upper Body
- Cobra Stretch: 3 sets of 30 seconds for the front body and lower back.
- Child’s Pose: 2 sets of 1 minute for full back relaxation.
- Thread the Needle: 2 sets of 30 seconds per side for upper back rotation.
- Doorway Chest Stretch: 2 sets of 30 seconds to fix "hunched" posture.
- Seated Spinal Twist: 3 sets of 20 seconds per side.

➡️Day 4: Active Flexibility Flow
- Stretch: 1 minute per side (lunge + thoracic twist).
- Dynamic Lunges with Twist: 2 sets of 10 reps.
- Deep Squat Hold: 1–2 minutes (opens ankles and hips).
- Wall Angels: 2 sets of 15 reps for shoulder mobility.
- Standing Side Stretch: 2 sets of 20 seconds per side.

➡️Day 5: Dynamic Stretching
- Sun Salutations: 3 rounds (gentle flow).
- Figure-Four Stretch: 2 minutes per side for glutes.
- Pancake Stretch: 2 sets of 1 minute (legs wide, leaning forward).
- Overhead Triceps Stretch: 30 seconds per side.
- Savasana (Corpse Pose): 3–5 minutes for final relaxation.
"""
    }

    text = plans.get(goal, "### 5-Day General Fitness Plan\nCustomized for your BMI and Age...")
    return [{"generated_text": text}]


# --- Helper Functions ---
def calculate_bmi(w, h_cm):
    if h_cm > 0:
        return round(w / ((h_cm / 100) ** 2), 2)
    return None


def get_bmi_cat(bmi):
    if bmi is None:
        return "Unknown"
    if bmi < 18.5:
        return "Underweight"
    if bmi < 25:
        return "Normal"
    if bmi < 30:
        return "Overweight"
    return "Obese"

# --- LOGIN SCREEN ---
if not st.session_state.authenticated:

    st.markdown("""
    <div class="hero-section">
        <h1 class="hero-title">🚀 FitPlan AI</h1>
        <p class="hero-subtitle">Your Ultimate AI-Powered Fitness Transformation</p>
    </div>
    """, unsafe_allow_html=True)

    col1, col2, col3 = st.columns([1,2,1])

    with col2:

        st.markdown('<div class="glass-card">', unsafe_allow_html=True)
        auth_mode = st.radio("Select Mode", ["Login","Sign Up"])

        # ---------------- SIGN UP ----------------
        if auth_mode == "Sign Up":

            name = st.text_input("👤 Name", key="signup_name")
            email = st.text_input("📧 Email", key="signup_email")
            password = st.text_input("🔑 Create 6-digit Password", type="password", key="signup_password")

            if st.button("📨 Send OTP"):

                if len(password) != 6:
                    st.error("Password must be 6 digits")
                    st.stop()

                if user_exists(email):
                    st.error("Email already registered")
                    st.stop()

                otp = generate_otp()

                st.session_state.signup_data = {
                    "name": name,
                    "email": email,
                    "password": password
                }

                st.session_state.otp = otp

                send_otp_email(email, otp)

                st.success("OTP sent to your email")

            if "otp" in st.session_state:

                otp_input = st.text_input("Enter OTP", key="signup_otp")

                if st.button("Verify OTP"):

                    if otp_input == st.session_state.otp:

                        data = st.session_state.signup_data

                        add_user(
                            data["name"],
                            data["email"],
                            data["password"]
                        )

                        st.success("Account created successfully!")

                        del st.session_state.otp

                    else:
                        st.error("Invalid OTP")

        # ---------------- LOGIN ----------------
        if auth_mode == "Login":

            email = st.text_input("📧 Email", key="login_email")
            password = st.text_input("🔑 Password", type="password", key="login_password")

            if st.button("Login"):

                user = verify_login(email, password)

                if user:

                    st.session_state.token = create_jwt(email)
                    st.session_state.authenticated = True

                    st.success("Login successful!")

                    st.rerun()

                else:
                    st.error("Invalid email or password")

        st.markdown('</div>', unsafe_allow_html=True)

# --- MAIN DASHBOARD ---
else:
    decoded = verify_jwt(st.session_state.token)
    if decoded:
        with st.sidebar:
            st.markdown(f"""
            <div class="glass-card" style="padding: 2rem; text-align: center;">
                <h3 style="color: #6366f1; margin-bottom: 1rem;">👋 Hello</h3>
                <p style="color: #64748b; font-size: 0.95rem;">{decoded['email']}</p>
            </div>
            """, unsafe_allow_html=True)
            page = st.radio("📱 Dashboard", ["💪 Profile", "🎯 My Plan"])
            if st.button("🚪 Sign Out"):
                for key in ['authenticated', 'token', 'otp', 'user_data']:
                    st.session_state.pop(key, None)
                st.rerun()

        if page == "💪 Profile":
            st.markdown("""
            <div class="hero-section">
                <h1 class="hero-title">🏋️ Create Profile</h1>
                <p class="hero-subtitle">Unlock your personalized AI workout plan</p>
            </div>
            """, unsafe_allow_html=True)

            st.markdown('<div class="glass-card">', unsafe_allow_html=True)
            st.markdown('<h2 style="color: #6366f1;">📊 Your Details</h2>', unsafe_allow_html=True)

            name = st.text_input("👤 Name")
            col1, col2, col3 = st.columns(3)
            with col1:
                gender = st.selectbox("Gender", ["Male", "Female", "Non-binary", "Prefer not to say"])
                height = st.number_input("📏 Height (cm)", value=170.0)
            with col2:
                age = st.number_input("🎂 Age", value=25)
                weight = st.number_input("⚖️ Weight (kg)", value=70.0)
            with col3:
                goal = st.selectbox("🎯 Goal", ["Build Muscle", "Weight Loss", "Strength Gain", "Abs Building", "Flexibility"])
                level = st.select_slider("⭐ Level", ["Beginner", "Intermediate", "Advanced"])

            equipment = st.multiselect(
                "🛠️ Equipment",
                ["Dumbbells", "Barbell", "Kettlebell", "Resistance Band", "Yoga Mat", "Full Gym", "Bodyweight"]
            )

            bmi = calculate_bmi(weight, height)

            if st.button("🚀 Generate My Plan", key="generate"):
                if name and equipment and bmi:
                    result = workout_generator("", goal)[0]["generated_text"]
                    st.session_state.user_data = {
                        "name": name,
                        "goal": goal,
                        "plan": result,
                        "bmi": bmi,
                        "bmi_status": get_bmi_cat(bmi)
                    }
                    st.success("✅ Plan ready! 👉 Check 'My Plan'")
                    st.rerun()
                else:
                    st.error("❌ Complete all fields")
            st.markdown('</div>', unsafe_allow_html=True)

        elif page == "🎯 My Plan":
            if not st.session_state.user_data:
                st.markdown("""
                <div class="glass-card" style="text-align: center; padding: 4rem;">
                    <h2 style="color: #f59e0b;">⚡ Ready to Start?</h2>
                    <p style="font-size: 1.3rem; color: #64748b;">
                        Create your profile first to unlock your AI-powered workout plan!
                    </p>
                </div>
                """, unsafe_allow_html=True)
            else:
                user = st.session_state.user_data
                st.markdown(f"""
                <div class="hero-section">
                    <h1 class="hero-title">Hey {user['name']}!</h1>
                    <p class="hero-subtitle">Your {user['goal']} Transformation Awaits ✨</p>
                </div>
                """, unsafe_allow_html=True)

                col1, col2 = st.columns(2)
                with col1:
                    st.markdown(f"""
                    <div class="glass-card" style="text-align:center;">
                        <h3 style="color:#10b981;margin-bottom:0.5rem;">BMI Score</h3>
                        <p style="font-size:2.4rem;font-weight:800;">{user['bmi']}</p>
                    </div>
                    """, unsafe_allow_html=True)
                with col2:
                    status = user['bmi_status']
                    color = "#10b981" if status == "Normal" else "#f59e0b"
                    st.markdown(f"""
                    <div class="glass-card" style="text-align:center;">
                        <h3 style="color:{color};margin-bottom:0.5rem;">Category</h3>
                        <p style="font-size:2rem;font-weight:700;">{status}</p>
                    </div>
                    """, unsafe_allow_html=True)

                st.markdown('<div class="glass-card">', unsafe_allow_html=True)
                st.markdown(f'<h2 style="color:#10b981; text-align:center;">🏋️ Your 5-Day {user["goal"]} Plan</h2>', unsafe_allow_html=True)
                st.markdown(f'<div class="plan-container">{user["plan"]}</div>', unsafe_allow_html=True)
                st.download_button(
                    "💾 Save Plan",
                    user["plan"],
                    f'{user["name"]}_{user["goal"]}_plan.txt',
                    use_container_width=True
                )
                st.markdown('</div>', unsafe_allow_html=True)

    else:
        st.markdown("""
        <div class="glass-card" style="text-align: center; padding: 3rem;">
            <h2 style="color: #ef4444;">⚠️ Session Expired</h2>
            <p style="color: #64748b;">Please login again</p>
        </div>
        """, unsafe_allow_html=True)
        if st.button("🔙 Back to Login"):
            st.rerun()