
from rapidfuzz import fuzz
from tqdm import tqdm

from src.Library import *
from src.helpers import *

import unicodedata

user_def_artist_map = {
    "donda" : "kanye west",
}

def normalize(s):
    return unicodedata.normalize("NFKC", s).replace("‐", "-")

def get_score(first: str, second: str) -> float:
    first_norm = normalize(first).lower()
    second_norm = normalize(second).lower()
    return float(first_norm == second_norm) * 100
    return fuzz.token_set_ratio(first, second)

def get_album_score(first: str, second: str) -> float:
    return get_score(first, second)

def get_track_score(first: str, second: str) -> float:
    return get_score(first, second)

def get_artist_score(first: str, second: str) -> float:
    if not first or not second:
        return 0.0

    first_lower = first.lower()
    second_lower = second.lower()

    for key, value in user_def_artist_map.items():
        key_lower = key.lower()
        value_lower = value.lower()
        if (first_lower == key_lower and second_lower == value_lower) or \
           (first_lower == value_lower and second_lower == key_lower):
            return 100.0

    return get_score(first, second)


def is_close_enough(first: float) -> bool:
    return first > 90

def generate_map(local_lib : Library, remote_lib : Library, output_dir : Path):

    artist_map : dict[str, tuple[str, float]] = {}
    album_map : dict[str, tuple[str, float]] = {}
    track_map : dict[str, tuple[str, float]] = {}

    with (tqdm(total=len(remote_lib.artists), desc='searching remote artists in local library') as process_bar):
        for remote_artist in remote_lib.artists:
            for local_artist in local_lib.artists:
                score = get_artist_score(remote_artist.title, local_artist.title)
                if not is_close_enough(score):
                    continue

                if remote_artist.id not in artist_map or score > artist_map[remote_artist.id][1]:
                    artist_map[remote_artist.id] = (local_artist.id, score)
                    remote_artist.local_id = local_artist.id

            process_bar.update(1)


    with tqdm(total=len(artist_map), desc='searching remote albums for each artist in local library') as process_bar:
        for remote_artist_id, (local_artist_id, _)  in artist_map.items():
            remote_artist = remote_lib.artist_map[remote_artist_id]
            local_artist = local_lib.artist_map[local_artist_id]

            for remote_album in remote_artist.albums:
                for local_album in local_artist.albums:
                    score = get_album_score(remote_album.title, local_album.title)
                    if not is_close_enough(score):
                        continue

                    if remote_album.id not in artist_map or score > album_map[remote_album.id][1]:
                        album_map[remote_album.id] = (local_album.id, score)
                        remote_album.local_id = local_album.id

            process_bar.update(1)

    with tqdm(total=len(album_map), desc='searching remote tracks for each album in local library') as process_bar:
        for remote_album_id, (local_album_id, _)  in album_map.items():
            remote_album = remote_lib.album_map[remote_album_id]
            local_album = local_lib.album_map[local_album_id]

            for remote_track in remote_album.tracks:
                for local_track in local_album.tracks:
                    score = get_track_score(remote_track.title, local_track.title)
                    if not is_close_enough(score):
                        continue

                    if remote_track.id not in artist_map or score > track_map[remote_track.id][1]:
                        track_map[remote_track.id] = (local_track.id, score)
                        remote_track.local_id = local_track.id

            process_bar.update(1)

    remote_lib.calc_resolved_percentage()

    out = {
        "artist_map" : { remote_track_id : local_track_id for remote_track_id, (local_track_id, _)  in artist_map.items() },
        "album_map" : { remote_track_id : local_track_id for remote_track_id, (local_track_id, _)  in album_map.items() },
        "track_map" : { remote_track_id : local_track_id for remote_track_id, (local_track_id, _)  in track_map.items() },
    }

    set_workdir(output_dir)
    save_data(out, "remote_to_local_map")
