#!/usr/bin/env python3

import os
import re
import logging
import argparse
from pathlib import Path
from difflib import SequenceMatcher

WORKDIR = Path("/data/navidrome/music/")
PLAYLIST_DIR = WORKDIR / "playlists"
MUSIC_DIR = WORKDIR / "music"
AUDIO_EXTS = {".mp3", ".flac", ".m4a", ".ogg", ".opus", ".aac", ".wav", ".ape"}
TRACE = 5
score_threshold = 0.6
folder_score_threshold = 0.6
logging.addLevelName(TRACE, "TRACE")
weight_jaccard = 0.6
weight_fuzzy = 1 - weight_jaccard


class CustomLogger(logging.Logger):
    def trace(self, msg: str, *args, **kwargs) -> None:
        if self.isEnabledFor(TRACE):
            self._log(TRACE, msg, args, **kwargs)


logging.setLoggerClass(CustomLogger)
logger: CustomLogger = logging.getLogger(__name__)  # type: ignore[assignment]


def words(text):
    return set(filter(None, re.split(r"[^a-z0-9]+", text.lower())))


def word_similarity(a: str, b: str) -> float:
    return SequenceMatcher(None, a, b).ratio()


def score_match_jaccard(search_words, candidate_words):
    intersection = search_words & candidate_words
    union = search_words | candidate_words

    if not union:
        return 0.0

    return len(intersection) / len(union)


def score_match_fuzzy(search_words, candidate_words):
    total = 0.0

    for sw in search_words:
        best = 0.0

        for cw in candidate_words:
            best = max(best, word_similarity(sw, cw))

        total += best

    #    logger.trace(f"Total: {total}/{len(search_words)}")
    return total / len(search_words)


def combined_score(score_fuzzy: float, score_jaccard: float) -> float:
    return score_fuzzy * weight_fuzzy + score_jaccard * weight_jaccard


def find_folder(directory: Path, folder: str) -> list[Path]:
    matches: list[tuple[Path, float]] = []
    search_words = words(folder)

    for candidate in directory.iterdir():
        if not candidate.is_dir():
            continue

        candidate_words = words(candidate.name)

        score_fuzzy = score_match_fuzzy(search_words, candidate_words)
        score_jaccard = score_match_jaccard(search_words, candidate_words)
        score = combined_score(score_fuzzy, score_jaccard)
        logger.trace(f"Candidate: {score:.2f}:{candidate_words}")
        if score > folder_score_threshold:
            matches.append((candidate, score))

    ret = [
        path
        for path, _ in sorted(
            matches,
            key=lambda x: x[1],
            reverse=True,
        )
    ]

    logger.debug(f"Find Folders Return: {ret}")
    return ret


def find_replacement(old_path: Path):
    p_old_path = Path(old_path)
    playlist_filename = p_old_path.stem
    # Find the second "music" directory
    music_idx = p_old_path.parts.index("music", p_old_path.parts.index("music") + 1)
    folder_count = len(p_old_path.parts) - music_idx - 2
    logger.debug(f"Folder Count Index: {folder_count}")
    try:
        artist = p_old_path.parts[len(WORKDIR.parts) + 1]  # /music/Artist/Album/File
        logger.debug(f"Artist: {artist}")
    except IndexError:
        logger.error(f"No Artist found for {p_old_path}")
        return None

    if artist == "Musik":
        artist_directories = [MUSIC_DIR / Path(playlist_filename.split(" - ")[0])]
    else:
        artist_directories = find_folder(MUSIC_DIR, artist)
    if not artist_directories:
        logger.error(f"Artist dir {artist_directories} dors not exist")
        return None

    logger.debug(f"Artist paths: {artist_directories}")
    album_directories = []
    subalbum_directories = []
    search_dir_list = artist_directories
    if folder_count >= 2:
        for artist_dir in artist_directories:
            try:
                album = p_old_path.parts[
                    len(WORKDIR.parts) + 2
                ]  # .../Artist/Album/File
                logger.debug(f"Album: {album}")
            except IndexError:
                logger.error(f"No album found for {p_old_path}")
                return None

            album_directories.extend(find_folder(artist_dir, album))
            if not album_directories:
                logger.warning(f"Album dir for '{album}' does not exist")
        if album_directories:
            logger.debug(f"Album path: {album_directories}")
            search_dir_list = album_directories

    if folder_count >= 3:
        for album_dir, score_fuzzy in album_directories:
            try:
                subalbum = p_old_path.parts[
                    len(WORKDIR.parts) + 3
                ]  # .../Artist/Album/File
                logger.debug(f"Sub-Album: {subalbum}")
            except IndexError:
                logger.error(f"No Sub-Album found for {p_old_path}")
                return None

            subalbum_directories.extend(find_folder(album_dir, subalbum))

            if not subalbum_directories:
                logger.error(f"Sub-Album dir for '{subalbum}' does not exist")
                return None

        if subalbum_directories:
            logger.debug(f"Sub-Album path: {subalbum_directories}")
            search_dir_list = subalbum_directories

    if not search_dir_list:
        logger.error(f"Search list empty")
        return None

    if (
        len(artist_directories) == 1
        and artist_directories[0].name.lower() in playlist_filename.lower()
    ):
        logger.debug(f"Replace {artist_directories[0].name} in {playlist_filename}")
        playlist_filename = re.sub(
            artist_directories[0].name, "", playlist_filename, flags=re.IGNORECASE
        )
        logger.debug(f"Filename minus Artist: {playlist_filename}")

    original = playlist_filename
    playlist_filename = re.sub(
        r"\s*[\(\[].*?[\)\]]",
        "",
        playlist_filename,
    )

    playlist_filename = re.sub(r"\s+", " ", playlist_filename).strip()

    if playlist_filename != original:
        logger.debug(f"Removed bracketed text: '{original}' -> '{playlist_filename}'")

    search_words = words(playlist_filename)
    best_match = None
    best_score = 0

    for search_dir in search_dir_list:
        for root, _, files in os.walk(search_dir):
            logger.debug(f"Search words {search_words}")
            logger.debug(f"Walk Directory: {root}")
            for filename in files:
                candidate = Path(root) / filename

                candidate_parts = candidate.stem.split(" - ")

                logger.trace(f"Candidate: {candidate}")
                if candidate.suffix.lower() not in AUDIO_EXTS:
                    continue

                separator_count = len(candidate.parent.name.split(" - ")) - 1
                if len(candidate_parts) == 4 + separator_count:
                    candidate_words = words(
                        " ".join(candidate_parts[2 + separator_count :])
                    )
                    filename_structure = "LidArr"
                else:
                    filename_structure = "Fallback"
                    candidate_words = words(candidate.stem)

                score_fuzzy = score_match_fuzzy(search_words, candidate_words)
                score_jaccard = score_match_jaccard(search_words, candidate_words)
                score = combined_score(score_fuzzy, score_jaccard)
                logger.trace(
                    f"{filename_structure} Words: {candidate_words} {score:.2f} / {best_score:.2f}, "
                    f"{weight_fuzzy:.2} * {score_fuzzy:.2f} + {weight_jaccard:.2} * {score_jaccard:.2f} "
                )
                if score > best_score:
                    best_score = score
                    best_match = candidate

    # Require all words to match
    if best_match and best_score > score_threshold:
        logger.debug(f"Best Match: {best_match}")

        return best_match

    logger.error(
        f"No match found with score greater {score_threshold:.2f}, best match: {best_score:.2f}:{best_match}"
    )
    return None


def main():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "path",
        nargs="?",
        default=PLAYLIST_DIR,
        help="Playlist file or directory",
    )

    parser.add_argument(
        "-ll",
        "--log-level",
        type=str.upper,
        default="INFO",
        choices=["TRACE", "DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
        help="Set log level",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Dry run only",
    )
    args = parser.parse_args()

    LOG_LEVELS = {
        "TRACE": TRACE,
        "DEBUG": logging.DEBUG,
        "INFO": logging.INFO,
        "WARNING": logging.WARNING,
        "ERROR": logging.ERROR,
        "CRITICAL": logging.CRITICAL,
    }

    logging.basicConfig(
        level=LOG_LEVELS[args.log_level.upper()],
        format="%(asctime)s %(levelname)s %(message)s",
    )
    playlist_path = Path(args.path)

    if playlist_path.is_file():
        playlists = [playlist_path.absolute()]
        logger.info(f"Edit single playlist: {playlists}")
    elif playlist_path.is_dir():
        playlists = [p for p in playlist_path.glob("*.m3u") if "_new.m3u" not in p.name]
        logger.info(f"Edit all playlists in {playlist_path.name}: {playlists}")
    else:
        logger.error(f"{playlist_path} does not exist")
        return

    for playlist in playlists:
        if "_new.m3u" in playlist.name:
            continue
        logger.info(f"\nProcessing {playlist}")

        changed = False
        output = []

        with open(playlist, "r", encoding="utf-8", errors="ignore") as f:
            lines = f.readlines()

        for line in lines:
            if not line or line.startswith("#"):
                # output.append(line)
                # Ignore EXINF line as they dont matter
                continue

            path = WORKDIR / Path(line.strip()).relative_to("/")

            if not path.is_absolute() or path.exists():
                output.append(line)
                continue

            logger.debug(f"Missing: {path}")
            replacement = find_replacement(path)
            if replacement:
                logger.info(f"{path.name} -> {replacement.name}")
                replacement = Path(
                    str(replacement).removeprefix("/data/navidrome/music")
                )
                output.append(str(replacement) + "\n")
                changed = True
            else:
                logger.error(f"  -> No match for {path}")
                output.append(line + "\n")
        if changed and not args.dry_run:
            backup_dir = playlist.parent.with_name(f"{playlist.parent.name}-bak")
            backup_playlist = backup_dir / playlist.name
            if not backup_playlist.exists():
                logger.info(f"Generate Backup playlist {backup_playlist}")
                backup_playlist.write_text(
                    playlist.read_text(encoding="utf-8"), encoding="utf-8"
                )

            with open(
                playlist.with_name(playlist.stem + ".m3u"),
                "w",
                encoding="utf-8",
                # playlist.with_name(playlist.stem + "_new.m3u"), "w", encoding="utf-8"
            ) as f:
                f.writelines(output)

            logger.info("Updated playlist")


if __name__ == "__main__":
    main()