import pandas as pd

# Read the CSV files
print("Reading music info csv ...")
tracks_df = pd.read_csv('data/music_info.csv')

print("Reading user listening history ...")
track_interactions_df = pd.read_csv('data/user_listening_history_10k.csv', nrows=1000)

# Merge the dataframes on 'track_id'
dataframe = pd.merge(tracks_df, track_interactions_df, on='track_id', how='left')

# Convert all NaN values to empty strings and all columns to string type
dataframe.fillna('', inplace=True)
dataframe = dataframe.astype(str)

# Group by 'user_id' and then create a list of dictionaries for each group
lookup_table = {user_id: group.drop('user_id', axis=1).to_dict('records')
                for user_id, group in dataframe.groupby('user_id')}

def get_users_with_track_interactions(ascending=False, limit=10):
    # Count the number of rows for each 'user_id'
    playcount_summary = track_interactions_df.groupby('user_id').size().reset_index(name='track_interactions')
    # Sort the DataFrame based on 'track_interactions', either ascending or descending
    playcount_summary.sort_values(by='track_interactions', ascending=ascending, inplace=True)
    # Limit the results if limit is specified
    if limit is not None:
        playcount_summary = playcount_summary.head(limit)
    # Convert the DataFrame to a list of dictionaries
    return playcount_summary.to_dict(orient='records')

def get_top_tracks_for_user(user_id: str, limit=20):
    # Retrieve the user's track list from the lookup table or an empty list if not found
    track_list = lookup_table.get(user_id, [])
    # Sort the track list by 'playcount' in descending order (assuming 'playcount' is stored as a string)
    sorted_tracks = sorted(track_list, key=lambda x: int(x['playcount']) if 'playcount' in x and x['playcount'].isdigit() else 0, reverse=True)
    # Apply the limit if specified
    if limit is not None:
        sorted_tracks = sorted_tracks[:limit]
    return sorted_tracks