import streamlit as st
import requests
import json
import os
import hashlib
from datetime import datetime
from dotenv import load_dotenv
from config import MODELS
load_dotenv(".env", override=True)

if 'authenticated' not in st.session_state:
    st.session_state.authenticated = False
if 'username' not in st.session_state:
    st.session_state.username = None

FIXED_USERNAME = os.getenv("USERNAME")
FIXED_PASSWORD_HASH = hashlib.sha256(os.getenv("PASSWORD").encode()).hexdigest()


st.set_page_config(
    page_title="AI Health Coach",
    page_icon="💪",
    layout="wide",
)

st.markdown("""
<style>
    .main-header {
        font-size: 2.5rem;
        color: #1E88E5;
        text-align: center;
        margin-bottom: 1rem;
    }
    .workout-card {
        background-color: #f0f2f6;
        border-radius: 10px;
        padding: 20px;
        margin-bottom: 15px;
    }
    .day-header {
        color: #1E88E5;
        font-size: 1.3rem;
        font-weight: bold;
    }
    .exercise-name {
        font-weight: bold;
        color: #333;
    }
    .exercise-detail {
        color: #555;
        margin-left: 10px;
    }
    .section-header {
        font-size: 1.5rem;
        color: #333;
        margin-top: 1rem;
        margin-bottom: 0.5rem;
    }
    .simple-login {
        max-width: 400px;
        margin: 100px auto;
        padding: 1.5rem;
        background-color: white;
        border-radius: 5px;
        box-shadow: 0 2px 10px rgba(0, 0, 0, 0.1);
    }
    .login-title {
        text-align: center;
        font-size: 1.5rem;
        margin-bottom: 1.5rem;
        color: #1E88E5;
    }
    .logout-btn {
        position: absolute;
        top: 10px;
        right: 10px;
    }
</style>
""", unsafe_allow_html=True)


def hash_password(password):
    return hashlib.sha256(password.encode()).hexdigest()

def verify_login(username, password):
    hashed_password = hash_password(password)
    return username == FIXED_USERNAME and hashed_password == FIXED_PASSWORD_HASH


def show_login_page():
    # Simplified login page with minimal elements
    col1, col2, col3 = st.columns([1, 2, 1])
    
    with col2:
        st.markdown("<div class='simple-login'>", unsafe_allow_html=True)
        st.markdown("<h3 class='login-title'>AI Health Coach</h3>", unsafe_allow_html=True)
        
        username = st.text_input("Username")
        password = st.text_input("Password", type="password")
        
        login_btn = st.button("Login", use_container_width=True)
        
        if login_btn:
            if verify_login(username, password):
                st.session_state.authenticated = True
                st.session_state.username = username
                st.success("Login successful!")
                st.rerun()
            else:
                st.error("Invalid username or password")
        
        st.markdown("</div>", unsafe_allow_html=True)
        
def show_main_app():
    st.markdown("<h1 class='main-header'>AI Zach Bouery </h1>", unsafe_allow_html=True)
    
    col1, col2, col3 = st.columns([1, 1, 1])
    with col3:
        st.markdown("<div style='text-align: right;'>", unsafe_allow_html=True)
        if st.button("Logout"):
            st.session_state.authenticated = False
            st.session_state.username = None
            st.experimental_rerun()
        st.markdown(f"Logged in as: **{st.session_state.username}**")
        st.markdown("</div>", unsafe_allow_html=True)
    
    if 'openrouter_config' not in st.session_state:
        st.session_state.openrouter_config = {"api_key": os.getenv('OpenrouterAPIKey')}
    if 'workout_plan' not in st.session_state:
        st.session_state.workout_plan = None
    if 'generation_time' not in st.session_state:
        st.session_state.generation_time = None
    if 'model_used' not in st.session_state:
        st.session_state.model_used = None
    
    with st.sidebar:
        st.header("Configuration")

        openrouter_api_key = os.getenv("OpenrouterAPIKey")
        print(openrouter_api_key)

        st.session_state.openrouter_config["api_key"] = openrouter_api_key
        
        st.subheader("Model Selection")
        model_options = list(MODELS.keys())
        selected_model = st.selectbox("Choose an LLM", model_options)
        
        st.subheader("User Information")
        user_age = st.number_input("Age", min_value=18, max_value=100, value=21)
        user_weight = st.number_input("Weight (kg)", min_value=30, max_value=100, value=70)
        user_height = st.number_input("Height (cm)", min_value=100, max_value=200, value=172)
        
        fitness_level = st.selectbox(
            "Fitness Level",
            options=["Beginner", "Intermediate", "Advanced"],
            index = 2
        )
        health_conditions = st.multiselect(
            "Health Conditions (if any)",
            ["None", "High Blood Pressure", "Diabetes", "Asthma", "Joint Pain", "Back Pain", "Heart Condition", "Other"]
        )
        
        time_available = st.text_input("Time Available Per Day (minutes)")
    
    col1, col2 = st.columns([1, 1])

    with col1:
        st.markdown("<div class='section-header'>Workout Goals</div>", unsafe_allow_html=True)
        primary_goal = st.selectbox(
            "Primary Goal",
            ["Weight Loss", "Muscle Building", "Cardiovascular Health", "Flexibility & Mobility", "General Fitness"]
        )
        secondary_goals = st.multiselect(
            "Secondary Goals (Optional)",
            ["Stress Reduction", "Better Sleep", "Increased Energy", "Improved Posture", "Core Strength", "Sport-Specific Training"]
        )
        equipment_available = st.multiselect(
            "Equipment Available",
            ["None (Bodyweight only)", "Dumbbells", "Resistance Bands", "Kettlebells", "Pull-up Bar", "Bench", "Full Gym Access"]
        )
        workout_days = st.text_input("Days Per Week", "3")
        workout_location = st.radio("Workout Location", ["Home", "Gym"])
        
    with col2:
        st.markdown("<div class='section-header'>Customize Prompt</div>", unsafe_allow_html=True)
        
        default_prompt = f"""
You are an expert fitness coach creating a personalized weekly workout plan. 
Primary goal: {primary_goal}
Secondary goals: {', '.join(secondary_goals) if secondary_goals else 'None'}
Fitness level: {fitness_level}
Age: {user_age}
Weight: {user_weight} kg
Height: {user_height} cm
Health conditions: {', '.join(health_conditions) if health_conditions else 'None'}
Time available: {time_available} minutes per session
Workout days: {workout_days} days per week
Equipment: {', '.join(equipment_available) if equipment_available else 'None'}
Location: {workout_location}

Create a detailed weekly workout plan with these requirements:
1. Include {workout_days} workout days with rest days appropriately spaced. 
2. Each workout should be completable within {time_available} minutes
3. Include general warm-up, then day-specific warmup
4. Include general cool-down recommendations and warm-up sets for each exercise
5. Provide specific exercises with sets, reps, rest periods, and recommended RPE 
6. Use RPE and RIR based approach
7. Use science based approach for exercise selection and progression
8. Do not recommend redundant exercises (ie, 2 exercises that target the same muscle, at the exact same angle and ROM)
9. Include progression tips
10. Do not give a specific day for core, recommend to do core exercises on every other day (max 3 exercises)
11. Format the response in a clear, organized way with days of the week as headers
12. Include a brief explanation of why this plan suits the user's goals and fitness level
13. Take into consideration that the user is at {fitness_level} fitness level

examples:
- for 6 days: Arnold Split OR PPL x2
- for 5 days: Push, Pull, Legs, Upper, Lower OR Upper, Lower, Upper, Lower, Weak Points
- for 4 days:  Upper, Lower, Upper, Lower
- for 3 days: Full Body x3
"""
        
        prompt_text = st.text_area("Customize Prompt", default_prompt, height=400)
        
        if st.button("Generate Workout Plan", type="primary"):
            if not st.session_state.openrouter_config["api_key"]:
                st.error("Please enter your OpenRouter API key.")
            else:
                with st.spinner("Generating your personalized workout plan..."):
                    try:
                        model_config = MODELS[selected_model]
                        
                        headers = {
                            "Authorization": f"Bearer {st.session_state.openrouter_config['api_key']}",
                            "Content-Type": "application/json"
                        }
                        
                        payload = {
                            "model": model_config["model_id"],
                            "messages": [
                                {"role": "system", "content": "You are an expert fitness coach specializing in creating personalized workout plans."},
                                {"role": "user", "content": prompt_text}
                            ],
                            "temperature": 1.8,
                            "max_tokens": 4000
                        }
                        
                        response = requests.post(
                            "https://openrouter.ai/api/v1/chat/completions",
                            headers=headers,
                            data=json.dumps(payload)
                        )
                                            
                        if response.status_code == 200:
                            response_data = response.json()
                        
                            
                            workout_plan = response_data["choices"][0]["message"]["content"]
                            
                            st.session_state.workout_plan = workout_plan
                            st.session_state.generation_time = datetime.now().strftime("%Y-%m-%d %H:%M")
                            st.session_state.model_used = selected_model
                            
                            st.success("Workout plan generated successfully!")
                        else:
                            error_details = f"Status Code: {response.status_code}"
                            try:
                                error_json = response.json()
                                error_details += f"\nError: {json.dumps(error_json, indent=2)}"
                            except:
                                error_details += f"\nResponse Text: {response.text}"
                            
                            st.error(f"API Error: {error_details}")
                    
                    except Exception as e:
                        st.error(f"An error occurred: {str(e)}")
                        import traceback
                        st.code(traceback.format_exc())


    if st.session_state.workout_plan:
        st.markdown("---")
        st.markdown("<div class='section-header'>Your Personalized Workout Plan</div>", unsafe_allow_html=True)
        
        col1, col2 = st.columns([1, 1])
        with col1:
            st.info(f"Generated on: {st.session_state.generation_time}")
        with col2:
            st.info(f"Model: {st.session_state.model_used}")
        st.markdown(st.session_state.workout_plan)

    st.markdown("---")
    st.markdown("© 2025 AI Health Coach | Powered by OpenRouter")


def main():
    if st.session_state.authenticated:
        show_main_app()
    else:
        show_login_page()

if __name__ == "__main__":
    main()