import re
import os
import sys
import json
from tqdm import tqdm
import polars as pl

from suno_utils.utils.text import read_jsonl


if __name__ == "__main__":

    dataset_name = "discogs"
    print(f"Loading {dataset_name} metas...")
    metas_filepath = f"/app/suno/tmp/raw_{dataset_name}_metas.jsonl"
    metas = read_jsonl(metas_filepath)

    # count the number of metas
    print(f"Number of metas: {len(metas)}")

    terms_set = [
        "audiophile",
        "remix",
        "cover",
        "live",
        "instrumental",
        "karaoke",
        "slowed + reverb",
        "official video",
        "official audio",
        "visualizer",
        "lyrics",
        "lyric video",
        "extended",
        "edit",
        "radio edit",
        "club mix",
        "bootleg",
        "rework",
        "re-edit",
        "reimagined",
        "demo",
        "acoustic",
        "unplugged",
        "reaction",
        "performance",
        "session",
        "mix",
        "dj set",
        "set",
        "tribute",
        "reprise",
        "duet",
        "mashup",
        "medley",
        "fan made",
        "fan edit",
        "loop",
        "3d audio",
        "8d audio",
        "360 audio",
        "slowed",
        "reverb",
        "sped up",
        "nightcore",
        "bass boosted",
        "bass-boosted",
        "hq",
        "hd",
        "4k",
        "vinyl",
        "vinyl rip",
        "cassette rip",
        "tape rip",
        "full album",
        "tracklist",
        "snippet",
        "teaser",
        "preview",
        "leak",
        "original mix",
        "alternate version",
        "studio version",
        "live version",
        "bedroom pop",
        "lofi",
        "lo-fi",
        "chill",
        "ambient",
        "synthwave",
        "phonk",
        "trap remix",
        "orchestral",
        "symphonic",
        "string version",
        "piano version",
        "guitar cover",
        "drum cover",
        "bass cover",
        "behind the scenes",
        "making of",
        "audio only",
        "concept video",
        "fancam",
        "practice",
        "dance practice",
        "line distribution",
        "tv performance",
        "award show",
        "web exclusive",
        "home recording",
        "bedroom recording",
        "raw take",
        "one take",
        "studio session",
        "jam session",
        "soundcheck",
        "backstage",
        "fan cam",
        "stage mix",
        "guitar",
        "electric guitar",
        "acoustic guitar",
        "bass",
        "bass guitar",
        "piano",
        "keyboard",
        "synth",
        "synthesizer",
        "drums",
        "percussion",
        "violin",
        "cello",
        "strings",
        "saxophone",
        "trumpet",
        "trombone",
        "flute",
        "clarinet",
        "harmonica",
        "banjo",
        "mandolin",
        "ukulele",
        "harp",
        "organ",
        "vocals",
        "vocal cover",
        "a cappella",
        "voice",
        "beatbox",
        "live performance",
        "remastered",
        "remaster",
        "orchestral version",
        "binaural",
        "unreleased",
        "music video",
        "360 video",
        "surround sound",
        "live from",
        "Tiny Desk",
        "MTV Unplugged",
        "concert version",
        "clean version",
        "explicit",
        "AMV",  # Anime Music Video
        "with lyrics",
        "without vocals",
        "no vocals",
        "trailer",
    ]

    terms_set = set(terms_set)
    print(f"Number of terms: {len(terms_set)}")

    # map from the meta_id to the terms
    terms_map = {}
    # Precompile regex patterns for better performance
    term_patterns = {
        term: re.compile(r"\b" + re.escape(term) + r"\b") for term in terms_set
    }

    pbar = tqdm(metas)
    for meta in pbar:
        if dataset_name == "discogs" or dataset_name == "discogs_subset":
            title_text = meta["title"].lower()  # Convert to lowercase once
        elif dataset_name == "genius":
            title_text = meta["youtube_title"].lower()
        else:
            raise ValueError(f"Dataset name {dataset_name} not supported")
        # Use precompiled patterns for faster matching
        found_terms = [
            term
            for term, pattern in term_patterns.items()
            if pattern.search(title_text)
        ]
        if found_terms:
            terms_map[meta["id"]] = found_terms
            pbar.set_postfix(found_terms=len(terms_map))

    # save the terms_map
    with open(
        f"/home/christian/code/christian/metadata/{dataset_name}_title_terms_map.json",
        "w",
    ) as f:
        json.dump(terms_map, f)