from flask import Flask, render_template, send_from_directory, request, redirect, url_for
import os
from collections import defaultdict

app = Flask(__name__)

# Path to your directory of ID folders
AUDIO_ROOT = '/Users/christiansteinmetz/Downloads/auk-clips-up-u-2'

# Cache for scanned data
id_to_models = defaultdict(list)
all_model_names = set()
id_list = []

def scan_audio_dirs():
    global id_to_models, all_model_names, id_list
    id_to_models.clear()
    all_model_names.clear()
    id_list.clear()

    for entry in os.scandir(AUDIO_ROOT):
        if entry.is_dir():
            id = entry.name
            id_list.append(id)
            mp3_files = [f for f in os.listdir(entry.path) if f.endswith('.mp3')]
            for f in mp3_files:
                # Extract model name
                if f.startswith(id + '_'):
                    model_name = f[len(id)+1:-4]  # strip ID_ and .mp3
                    id_to_models[id].append(model_name)
                    all_model_names.add(model_name)
    id_list.sort()

    # print a summary of the files 
    print("--------------------------------")
    print(f"Found {len(id_list)} ids")
    print(f"Found {len(all_model_names)} model names")
    print(f"Found {len(id_to_models)} id to models")
    print("--------------------------------")
# Initial scan
scan_audio_dirs()

@app.route('/')
def index():
    return render_template('index.html', model_names=sorted(all_model_names))

@app.route('/compare')
def compare():
    selected_models = request.args.getlist('models')
    if not selected_models:
        # You can redirect back to index with a message if you prefer
        return "Please select at least one model on the previous page.", 400

    # Build a filtered list of IDs where all selected models exist
    required = set(selected_models)
    filtered_ids = [
        id_ for id_ in id_list
        if required.issubset(set(id_to_models.get(id_, [])))
    ]

    if not filtered_ids:
        # Nothing to show with this selection
        return render_template(
            'compare.html',
            current_id=None,
            model_files=[],
            selected_models=selected_models,
            page=0,
            total_pages=0
        )

    page = int(request.args.get('page', 0))
    if page < 0 or page >= len(filtered_ids):
        return "Invalid page", 404

    current_id = filtered_ids[page]

    # Build the files list (should exist by construction, but double-check)
    model_files = []
    for model in selected_models:
        filename = f"{current_id}_{model}.mp3"
        file_path = os.path.join(AUDIO_ROOT, current_id, filename)
        if os.path.exists(file_path):
            model_files.append({'model': model, 'filename': filename})

    return render_template(
        'compare.html',
        current_id=current_id,
        model_files=model_files,
        selected_models=selected_models,
        page=page,
        total_pages=len(filtered_ids)
    )

@app.route('/audio/<id>/<filename>')
def serve_audio(id, filename):
    return send_from_directory(os.path.join(AUDIO_ROOT, id), filename)

if __name__ == '__main__':
    app.run(debug=True, port=8080)