import streamlit as st
import sqlite3
import hashlib
from transformers import pipeline

# -----------------------------
# LOAD AI MODEL (Milestone 2)
# -----------------------------
@st.cache_resource
def load_model():
    generator = pipeline(
        "text-generation",
        model="google/flan-t5-base"
    )
    return generator

generator = load_model()

# -----------------------------
# DATABASE SETUP
# -----------------------------
conn = sqlite3.connect("users.db", check_same_thread=False)
c = conn.cursor()

c.execute("""
CREATE TABLE IF NOT EXISTS users (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    username TEXT UNIQUE,
    email TEXT,
    password TEXT,
    age INTEGER
)
""")
conn.commit()

# -----------------------------
# PASSWORD HASH
# -----------------------------
def hash_password(password):
    return hashlib.sha256(password.encode()).hexdigest()

# -----------------------------
# PAGE CONFIG
# -----------------------------
st.set_page_config(page_title="FitPlan AI", layout="wide")

# -----------------------------
# SESSION STATE
# -----------------------------
if "logged_in" not in st.session_state:
    st.session_state.logged_in = False

if "username" not in st.session_state:
    st.session_state.username = None


# -----------------------------
# SIGNUP
# -----------------------------
def signup():
    st.title("📝 Create Account")

    username = st.text_input("Username")
    email = st.text_input("Email")
    password = st.text_input("Password", type="password")

    if st.button("Sign Up"):
        if username and email and password:
            hashed_pw = hash_password(password)
            try:
                c.execute("INSERT INTO users (username, email, password) VALUES (?, ?, ?)",
                          (username, email, hashed_pw))
                conn.commit()
                st.success("Account Created Successfully! Please Login.")
            except:
                st.error("Username already exists!")
        else:
            st.warning("Please fill all fields.")


# -----------------------------
# LOGIN
# -----------------------------
def login():
    st.title("🔐 Login")

    username = st.text_input("Username")
    password = st.text_input("Password", type="password")

    if st.button("Login"):
        hashed_pw = hash_password(password)

        c.execute("SELECT * FROM users WHERE username=?", (username,))
        user = c.fetchone()

        if user and user[3] == hashed_pw:
            st.session_state.logged_in = True
            st.session_state.username = username
            st.success("Login Successful!")
            st.rerun()
        else:
            st.error("Invalid Username or Password")


# -----------------------------
# PROFILE + WORKOUT INPUT
# -----------------------------
def profile_page():
    st.title("🏋️ FitPlan AI - Personalized Workout Generator")

    st.write(f"Welcome, **{st.session_state.username}**")

    # Get age
    c.execute("SELECT age FROM users WHERE username=?",
              (st.session_state.username,))
    data = c.fetchone()
    existing_age = data[0] if data[0] else 18

    age = st.number_input("Age", 10, 100, existing_age)

    # Workout Inputs (Milestone 1)
    goal = st.selectbox("Fitness Goal",
                        ["Build Muscle", "Lose Weight", "Improve Cardio", "Flexibility"])

    level = st.selectbox("Fitness Level",
                         ["Beginner", "Intermediate", "Advanced"])

    equipment = st.selectbox("Available Equipment",
                             ["No Equipment", "Dumbbells", "Full Gym"])

    if st.button("Save Profile"):
        c.execute("UPDATE users SET age=? WHERE username=?",
                  (age, st.session_state.username))
        conn.commit()
        st.success("Profile Updated!")

    # -----------------------------
    # GENERATE WORKOUT PLAN
    # -----------------------------
    if st.button("Generate Workout Plan"):

        with st.spinner("Generating your personalized workout plan..."):

            prompt = f"""
            Generate a structured 5-day workout plan.
            Age: {age}
            Goal: {goal}
            Fitness Level: {level}
            Equipment: {equipment}

            Format clearly like:

            Day 1:
            Warm-up:
            Exercises:
            Sets/Reps:
            Rest:

            Continue for 5 days.
            """

            result = generator(prompt, max_length=500, do_sample=True)
            plan = result[0]["generated_text"]

            st.subheader("📅 Your 5-Day Personalized Plan")
            st.write(plan)

    if st.button("Logout"):
        st.session_state.logged_in = False
        st.session_state.username = None
        st.rerun()


# -----------------------------
# NAVIGATION
# -----------------------------
if not st.session_state.logged_in:
    menu = st.sidebar.selectbox("Menu", ["Login", "Sign Up"])

    if menu == "Login":
        login()
    else:
        signup()
else:
    profile_page()