import matplotlib
# allow running headless
matplotlib.use('agg')

from flask import Flask, render_template, g, abort, redirect, session, request, url_for
from markupsafe import escape
import spotipy
from typing import Optional
from dotenv import load_dotenv
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure
import numpy as np
import io
import os
import base64

load_dotenv()

app = Flask(__name__)
app.secret_key = os.environ.get('SECRET_KEY')

def get_spotify() -> Optional[spotipy.Spotify]:
    cache_handler = spotipy.cache_handler.FlaskSessionCacheHandler(session)
    auth_manager = spotipy.oauth2.SpotifyOAuth(scope='playlist-read-private', cache_handler=cache_handler, show_dialog=True)
    if not auth_manager.validate_token(cache_handler.get_cached_token()):
        return None
    spotify = spotipy.Spotify(auth_manager=auth_manager)
    return spotify

@app.route('/playlist/<id>')
def playlist(id: str):
    spotify = get_spotify()
    if not spotify:
        return redirect(url_for('login', return_playlist=id))

    try:
        playlist = spotify.playlist(playlist_id=id)

        tracks = []
        result = spotify.playlist_items(playlist_id=id)
        tracks.extend(result['items'])

        # if playlist is larger than 100 songs, continue loading it until end
        while result['next']:
            result = spotify.next(result)
            tracks.extend(result['items'])

        artists = {}
        for track in tracks:
            if 'artists' in track['track']:
                for artist in track['track']['artists']:
                    if artists.get(artist['uri'], None):
                        artists[artist['uri']]['count'] += 1
                    else:
                        artists[artist['uri']] = {
                            'count': 1,
                            'name': artist['name'],
                        }
        artists = list(map(lambda artist: artist, artists.values()))
        artists.sort(key=lambda artist: artist['name'])
        artists.sort(key=lambda artist: -artist['count'])

        counts = list(map(lambda artist: artist['count'], artists))
        labels = list(map(lambda artist: artist['name'] + f' ({round(artist["count"] / len(tracks) * 100)}%)', artists))

        fig = Figure()
        axis = fig.add_subplot(1, 1, 1)

        axis.pie(counts, labels=labels, rotatelabels=True, explode=[.2] + ([0] * (len(artists) - 1)))
        data = io.BytesIO()
        FigureCanvas(fig).print_figure(data, format='png', bbox_inches='tight')
        pie = base64.b64encode(data.getvalue()).decode('ascii')
    except spotipy.SpotifyException as e:
        if e.code == 404:
            return abort(404)
        else:
            return abort(500)

    return render_template('playlist.html', playlist=playlist, tracks=tracks, track_count=len(tracks), artists=artists, pie=pie)

@app.route('/login')
def login():
    cache_handler = spotipy.cache_handler.FlaskSessionCacheHandler(session)
    auth_manager = spotipy.oauth2.SpotifyOAuth(scope='playlist-read-private', cache_handler=cache_handler, show_dialog=True)
    if request.args.get('code'):
        auth_manager.get_access_token(request.args.get('code'))
        if session.get('return_playlist'):
            return redirect(url_for('playlist', id=session.get('return_playlist')))
        else:
            return redirect('/')

    if not auth_manager.validate_token(cache_handler.get_cached_token()):
        if request.args.get('return_playlist'):
            session['return_playlist'] = request.args.get('return_playlist')
        return redirect(auth_manager.get_authorize_url())

    return redirect('/')

@app.route('/logout')
def logout():
    session.pop("token_info", None)
    return redirect('/')

@app.route('/')
def index():
    spotify = get_spotify()
    if spotify:
        playlists = []
        result = spotify.current_user_playlists()
        playlists.extend(result['items'])

        while result['next']:
            result = spotify.next(result)
            playlists.extend(result['items'])
        return render_template('index.html', playlists=playlists)
    else:
        return render_template('index-login.html')

def main():
    app.run(host='0.0.0.0', debug=False)

if __name__ == '__main__':
    main()