from flask import Blueprint, render_template, redirect, url_for, request, current_app
from jinja2_fragments import render_block
from flask_htmx import HTMX
from flask_login import login_required, current_user
from extensions import db
from decorators import validate_workout, validate_topset, require_ownership, validate_person
from utils import convert_str_to_date
from collections import defaultdict # Import defaultdict

workout_bp = Blueprint('workout', __name__)
htmx = HTMX()

# --- Database Helper Function (Moved from features/workout.py) ---

def _get_workout_view_model(person_id, workout_id):
    """Fetches and formats workout data for display."""
    query = """
    SELECT
        w.workout_id,
        w.person_id,
        p.name AS person_name,
        w.start_date,
        w.note,
        t.topset_id,
        e.exercise_id,
        e.name AS exercise_name,
        t.repetitions,
        t.weight,
        tag.tag_id,
        tag.name AS tag_name,
        tag.filter as tag_filter,
        -- Check if the tag is associated with *this specific workout*
        EXISTS (
            SELECT 1 FROM workout_tag wt_check
            WHERE wt_check.workout_id = w.workout_id AND wt_check.tag_id = tag.tag_id
        ) AS is_selected
    FROM
        workout w
    LEFT JOIN
        person p ON w.person_id = p.person_id
    LEFT JOIN
        topset t ON w.workout_id = t.workout_id
    LEFT JOIN
        exercise e ON t.exercise_id = e.exercise_id
    -- Cross Join with tags relevant to the person to get all possible tags
    CROSS JOIN
        tag
    WHERE
        w.person_id = %s AND w.workout_id = %s
        -- Ensure we only get tags associated with the person or general tags (person_id IS NULL)
        AND (tag.person_id = w.person_id OR tag.person_id IS NULL)
    ORDER BY t.topset_id ASC, tag.name ASC;
    """
    # Note: The original query had a potential issue joining tags.
    # The revised query uses a CROSS JOIN with person-specific tags
    # and then checks association using EXISTS for the 'is_selected' flag.
    # This ensures all relevant tags for the person are listed,
    # and correctly marks those applied to this workout.

    data = db.execute(query, [person_id, workout_id])

    if not data:
        # If workout exists but has no topsets/tags, fetch basic info
        basic_info = db.execute(
            """SELECT w.workout_id, w.person_id, p.name as person_name, w.start_date, w.note
               FROM workout w JOIN person p ON w.person_id = p.person_id
               WHERE w.person_id = %s AND w.workout_id = %s""",
            [person_id, workout_id], one=True
        )
        if not basic_info:
             return {"error": "Workout not found"}, 404 # Truly not found
        # Workout exists but is empty
        return {
            "workout_id": basic_info["workout_id"],
            "person_id": basic_info["person_id"],
            "person_name": basic_info["person_name"],
            "start_date": basic_info["start_date"],
            "note": basic_info["note"],
            "tags": [], # Fetch all person tags separately if needed
            "top_sets": [],
        }


    # Initialize workout dictionary using the first row for base details
    workout_data = {
        "workout_id": data[0]["workout_id"],
        "person_id": data[0]["person_id"],
        "person_name": data[0]["person_name"],
        "start_date": data[0]["start_date"],
        "note": data[0]["note"],
        "tags": [],
        "top_sets": [],
    }

    # Use sets to track unique IDs efficiently
    tag_ids_seen = set()
    topset_ids_seen = set()

    for row in data:
        # Process Tags
        tag_id = row.get("tag_id")
        if tag_id is not None and tag_id not in tag_ids_seen:
            workout_data["tags"].append({
                "tag_id": tag_id,
                "tag_name": row.get("tag_name"),
                "tag_filter": row.get("tag_filter"),
                "person_id": person_id, # Assuming tags are person-specific or global
                "is_selected": row.get("is_selected", False)
            })
            tag_ids_seen.add(tag_id)

        # Process Topsets
        topset_id = row.get("topset_id")
        if topset_id is not None and topset_id not in topset_ids_seen:
            workout_data["top_sets"].append({
                "topset_id": topset_id,
                "exercise_id": row.get("exercise_id"),
                "exercise_name": row.get("exercise_name"),
                "repetitions": row.get("repetitions"),
                "weight": row.get("weight")
            })
            topset_ids_seen.add(topset_id)

    # Sort tags alphabetically by name for consistent display
    workout_data["tags"].sort(key=lambda x: x.get('tag_name', ''))

    # Add multi-category breakdowns
    workout_data["muscle_distribution"] = db.exercises.get_workout_attribute_distribution(workout_id, 'Muscle Group')
    workout_data["equipment_distribution"] = db.exercises.get_workout_attribute_distribution(workout_id, 'Machine vs Free Weight')
    workout_data["movement_distribution"] = db.exercises.get_workout_attribute_distribution(workout_id, 'Compound vs Isolation')

    return workout_data


# --- Routes ---

@workout_bp.route("/person/<int:person_id>/workout", methods=['POST'])
@login_required
@validate_person
@require_ownership
def create_workout(person_id):
    new_workout_id = db.create_workout(person_id)
    db.activityRequest.log(current_user.id, 'CREATE_WORKOUT', 'workout', new_workout_id, f"Created workout for person_id: {person_id}")
    # Use the local helper function to get the view model
    view_model = _get_workout_view_model(person_id, new_workout_id)
    if "error" in view_model: # Handle case where workout creation might fail or is empty
         # Decide on appropriate error handling, maybe redirect or show error message
         return redirect(url_for('calendar.get_calendar', person_id=person_id))
    return render_block(current_app.jinja_env, 'workout.html', 'content', **view_model)

@workout_bp.route("/person/<int:person_id>/workout/<int:workout_id>/delete", methods=['GET'])
@login_required
@validate_workout
@require_ownership
def delete_workout(person_id, workout_id):
    db.delete_workout(workout_id)
    db.activityRequest.log(current_user.id, 'DELETE_WORKOUT', 'workout', workout_id, f"Deleted workout: {workout_id}")
    return redirect(url_for('calendar.get_calendar', person_id=person_id))

@workout_bp.route("/person/<int:person_id>/workout/<int:workout_id>/start_date_edit_form", methods=['GET'])
@login_required
@validate_workout
@require_ownership
def get_workout_start_date_edit_form(person_id, workout_id):
    # Fetch only the necessary data (start_date)
    workout = db.execute("SELECT start_date FROM workout WHERE workout_id = %s", [workout_id], one=True)
    start_date = workout.get('start_date') if workout else None
    return render_template('partials/start_date.html', person_id=person_id, workout_id=workout_id, start_date=start_date, is_edit=True)

@workout_bp.route("/person/<int:person_id>/workout/<int:workout_id>/start_date", methods=['PUT'])
@login_required
@validate_workout
@require_ownership
def update_workout_start_date(person_id, workout_id):
    new_start_date_str = request.form.get('start-date')
    db.update_workout_start_date(workout_id, new_start_date_str)
    db.activityRequest.log(current_user.id, 'UPDATE_WORKOUT_START_DATE', 'workout', workout_id, f"Updated start date to {new_start_date_str}")
    # Convert string back to date for rendering the partial
    new_start_date = convert_str_to_date(new_start_date_str, '%Y-%m-%d')
    return render_template('partials/start_date.html', person_id=person_id, workout_id=workout_id, start_date=new_start_date)

@workout_bp.route("/person/<int:person_id>/workout/<int:workout_id>/start_date", methods=['GET'])
@validate_workout
def get_workout_start_date(person_id, workout_id):
    # Fetch only the necessary data (start_date)
    workout = db.execute("SELECT start_date FROM workout WHERE workout_id = %s", [workout_id], one=True)
    start_date = workout.get('start_date') if workout else None
    return render_template('partials/start_date.html', person_id=person_id, workout_id=workout_id, start_date=start_date)

@workout_bp.route("/person/<int:person_id>/workout/<int:workout_id>/distribution", methods=['GET'])
def get_workout_distribution(person_id, workout_id):
    category = request.args.get('category', 'Muscle Group')
    distribution = db.exercises.get_workout_attribute_distribution(workout_id, category)
    return render_template('partials/workout_breakdown.html', 
                           person_id=person_id, 
                           workout_id=workout_id, 
                           distribution=distribution,
                           category_name=category)

@workout_bp.route("/person/<int:person_id>/workout/<int:workout_id>/topset/<int:topset_id>/achievements", methods=['GET'])
@validate_topset
def get_topset_achievements(person_id, workout_id, topset_id):
    achievements = db.get_topset_achievements(topset_id)
    return render_template('partials/achievement_badges.html', achievements=achievements)

@workout_bp.route("/person/<int:person_id>/workout/<int:workout_id>/topset/<int:topset_id>", methods=['GET'])
@validate_topset
def get_topset(person_id, workout_id, topset_id):
    topset = db.get_topset(topset_id) # Keep using db.py for simple gets for now
    return render_template('partials/topset.html', person_id=person_id, workout_id=workout_id, topset_id=topset_id, exercise_id=topset.get('exercise_id'), exercise_name=topset.get('exercise_name'), repetitions=topset.get('repetitions'), weight=topset.get('weight'))

@workout_bp.route("/person/<int:person_id>/workout/<int:workout_id>/topset/<int:topset_id>/edit_form", methods=['GET'])
@login_required
@validate_topset
@require_ownership
def get_topset_edit_form(person_id, workout_id, topset_id):
    exercises = db.get_all_exercises()
    topset = db.get_topset(topset_id)
    return render_template('partials/topset.html', person_id=person_id, workout_id=workout_id, topset_id=topset_id, exercises=exercises, exercise_id=topset.get('exercise_id'), exercise_name=topset.get('exercise_name'), repetitions=topset.get('repetitions'), weight=topset.get('weight'), is_edit=True)

@workout_bp.route("/person/<int:person_id>/workout/<int:workout_id>/topset", methods=['POST'])
@login_required
@validate_workout
@require_ownership
def create_topset(person_id, workout_id):
    exercise_id = request.form.get("exercise_id")
    repetitions = request.form.get("repetitions")
    weight = request.form.get("weight")
    new_topset_id = db.create_topset(workout_id, exercise_id, repetitions, weight)
    exercise = db.get_exercise(exercise_id)
    db.activityRequest.log(current_user.id, 'ADD_SET', 'topset', new_topset_id, f"Added set: {repetitions} x {weight}kg {exercise['name']} in workout {workout_id}")
    return render_template('partials/topset.html', person_id=person_id, workout_id=workout_id, topset_id=new_topset_id, exercise_id=exercise_id, exercise_name=exercise.get('name'), repetitions=repetitions, weight=weight), 200, {"HX-Trigger": "topsetAdded"}

@workout_bp.route("/person/<int:person_id>/workout/<int:workout_id>/topset/<int:topset_id>", methods=['PUT'])
@login_required
@validate_workout
@require_ownership
def update_topset(person_id, workout_id, topset_id):
    exercise_id = request.form.get("exercise_id")
    repetitions = request.form.get("repetitions")
    weight = request.form.get("weight")
    db.update_topset(exercise_id, repetitions, weight, topset_id)
    exercise = db.get_exercise(exercise_id)
    db.activityRequest.log(current_user.id, 'UPDATE_SET', 'topset', topset_id, f"Updated set {topset_id}: {repetitions} x {weight}kg {exercise['name']}")
    return render_template('partials/topset.html', person_id=person_id, workout_id=workout_id, topset_id=topset_id, exercise_id=exercise_id, exercise_name=exercise.get('name'), repetitions=repetitions, weight=weight)

@workout_bp.route("/person/<int:person_id>/workout/<int:workout_id>/topset/<int:topset_id>/delete", methods=['DELETE'])
@login_required
@validate_topset
@require_ownership
def delete_topset(person_id, workout_id, topset_id):
    topset = db.get_topset(topset_id)
    db.delete_topset(topset_id)
    db.activityRequest.log(current_user.id, 'DELETE_SET', 'topset', topset_id, f"Deleted set {topset_id}: {topset['exercise_name']} in workout {workout_id}")
    return ""

@workout_bp.route("/person/<int:person_id>/workout/<int:workout_id>/exercise/most_recent_topset_for_exercise", methods=['GET'])
def get_most_recent_topset_for_exercise(person_id, workout_id):
    exercise_id = request.args.get('exercise_id', type=int)
    exercises = db.get_all_exercises()
    if not exercise_id:
        return render_template('partials/new_set_form.html', person_id=person_id, workout_id=workout_id, exercises=exercises)
    topset = db.get_most_recent_topset_for_exercise(person_id, exercise_id) # Keep using db.py for now
    if not topset:
        exercise = db.execute("select name from exercise where exercise_id=%s", [exercise_id], one=True)
        return render_template('partials/new_set_form.html', person_id=person_id, workout_id=workout_id, exercises=exercises, exercise_id=exercise_id, exercise_name=exercise['name'])
    (repetitions, weight, exercise_name) = topset
    return render_template('partials/new_set_form.html', person_id=person_id, workout_id=workout_id, exercise_id=exercise_id, exercise_name=exercise_name, repetitions=repetitions, weight=weight)

@workout_bp.route("/person/<int:person_id>/workout/<int:workout_id>", methods=['GET'])
def show_workout(person_id, workout_id):
    # Use the local helper function to get the view model
    view_model = _get_workout_view_model(person_id, workout_id)
    if "error" in view_model:
         # Decide on appropriate error handling
         return redirect(url_for('calendar.get_calendar', person_id=person_id)) # Redirect to calendar on error
    if htmx:
        return render_block(current_app.jinja_env, 'workout.html', 'content', **view_model)
    return render_template('workout.html', **view_model)