from suno_utils.utils.text import read_jsonl, write_jsonl
from tqdm import tqdm
import os
import re
from collections import defaultdict

FOLDER = "/app2/suno/data/sara/imslp_filter/"
THRESHOLD = 1.0
OUTPUT_PATH = os.path.join(FOLDER, f"metas_raw_labeled_dupes_thresh_{THRESHOLD}.jsonl")
ARTIST_KEY = "composer"  # or artists

artist_full = read_jsonl(os.path.join(FOLDER, "metas_raw.jsonl"), progress=True)
metas_v0 = read_jsonl(os.path.join(FOLDER, "metas_v0.jsonl"), progress=True)
artist_full_clean_copy = artist_full.copy()

PAREN_PATTERN = re.compile(r"\([^)]*\)")
BRACKET_PATTERN = re.compile(r"\[[^\]]*\]")


def preprocess_text(text: str) -> str:
    """
    Preprocess text (title or artist name) for better comparison.

    Args:
        text: The original text

    Returns:
        Preprocessed text
    """
    if not text:
        return ""

    # Convert to lowercase
    processed = text.lower()

    # Helper function to clean song titles (remove text in brackets)
    def clean_title(title):
        # Remove text in parentheses and square brackets
        clean = PAREN_PATTERN.sub("", title)
        clean = BRACKET_PATTERN.sub("", clean)
        return clean.strip().lower()

    # processed = clean_title(processed)

    # Remove special characters (but preserve numbers)
    processed = re.sub(r"[!?:;#&+=%*\-_,.]", "", processed)

    # Remove extra whitespace
    processed = re.sub(r"\s+", " ", processed)

    return processed.strip()


def find_song_duplicates(song_entries, min_substring_length=4):
    potential_duplicates = []

    # Preprocess all song titles - O(n) operation
    processed_entries = []
    exact_match_dict = defaultdict(list)

    for id, title in song_entries:
        stripped_title = title.strip()
        cleaned_title = preprocess_text(stripped_title)
        # Skip titles that are too short after cleaning
        if len(cleaned_title) < min_substring_length:
            continue

        processed_entries.append((id, stripped_title, cleaned_title))

        # Add to dictionary for exact match detection
        exact_match_dict[cleaned_title].append((id, stripped_title))

    # First check for exact matches - O(n) operation
    for cleaned_title, entries in exact_match_dict.items():
        if len(entries) > 1:
            # Found exact matches
            for i in range(len(entries)):
                for j in range(i + 1, len(entries)):
                    potential_duplicates.append((entries[i], entries[j]))

    # Sort processed entries by cleaned title length to enable early termination
    processed_entries.sort(key=lambda x: len(x[2]))

    if THRESHOLD == 1.0:  # exact matches only
        return potential_duplicates

    # Then check for substring matches
    for i in range(len(processed_entries)):
        id1, title1, clean1 = processed_entries[i]

        # Only need to check against longer strings (shorter strings can't contain this one)
        for j in range(i + 1, len(processed_entries)):
            id2, title2, clean2 = processed_entries[j]

            # Skip exact matches (we already handled them)
            if clean1 == clean2:
                continue

            # Skip if the length ratio is too small (prevents false positives)
            if len(clean1) / len(clean2) < THRESHOLD:
                continue

            # Check if one is a substring of the other
            if clean1 in clean2 or clean2 in clean1:
                potential_duplicates.append(((id1, title1), (id2, title2)))

    return potential_duplicates


ids_to_title = {}
id_to_data = {}
for song_info in artist_full:
    if ARTIST_KEY not in song_info or "title" not in song_info:
        continue

    artist_info = song_info[ARTIST_KEY]
    if isinstance(artist_info, str):
        artist_ids = [artist_info]
    else:
        artist_ids = [x["id"] for x in song_info[ARTIST_KEY]]

    title = song_info["title"]
    id = song_info["id"]
    id_to_data[id] = song_info
    for a_id in artist_ids:
        if a_id not in ids_to_title:
            ids_to_title[a_id] = []
        ids_to_title[a_id].append((id, title))

for id, data in id_to_data.items():
    data["duplicate"] = False

num_tracks = 0
num_dupes = 0
dupes = []
for a_id, titles in tqdm(ids_to_title.items()):
    dupes_maybe = find_song_duplicates(titles)
    # print(a_id, titles)
    if len(dupes_maybe) > 0:
        num_tracks += len(titles)
        num_dupes += len(dupes_maybe)
        for dupe_source, dupe in dupes_maybe:
            # print(dupe_source, dupe)
            # print(dupe_source[1], dupe[1])
            dupes.append((dupe_source[1], dupe[1]))

            source = dupe_source[0]
            if "duplicate_source" in id_to_data[dupe_source[0]]:
                source = id_to_data[dupe_source[0]]["duplicate_source"]

            id_to_data[dupe[0]]["duplicate"] = True
            id_to_data[dupe[0]]["duplicate_source"] = source
            id_to_data[dupe_source[0]]["duplicate"] = True
            id_to_data[dupe_source[0]]["duplicate_source"] = source
print(f"{num_tracks} tracks, {num_dupes} dupe pairs")

cleaned_metas = []
filtered_metas = []
for meta in artist_full_clean_copy:
    id = meta["id"]
    if id in id_to_data:
        is_dupe = id_to_data[id]["duplicate"]
        meta["is_artist_duplicate"] = is_dupe
        if is_dupe:
            meta["duplicate_source"] = id_to_data[id]["duplicate_source"]
        if not is_dupe:
            filtered_metas.append(meta)
    else:
        filtered_metas.append(meta)
    cleaned_metas.append(meta)

num_duplicates = len(metas_v0) - len(filtered_metas)
percent_dupe = num_duplicates / len(metas_v0) * 100
print(f"Found {num_duplicates} duplicate titles out of {len(metas_v0)} tracks, {percent_dupe}%")

print(f"Writing output to {OUTPUT_PATH}")
write_jsonl(cleaned_metas, OUTPUT_PATH)