import os
import csv
import time
import json
import numpy as np

from tqdm import tqdm
from rapidfuzz import fuzz
from suno_utils.utils.text import read_jsonl, write_jsonl
from concurrent.futures import ProcessPoolExecutor, as_completed


# read csv file into a list of dictionaries
def read_csv(file_path):
    with open(file_path, mode="r") as file:
        reader = csv.DictReader(file)
        return [row for row in reader]


def filter_meta_batch(metas, artist_set):
    """Process a batch of metas."""
    filtered = []
    for meta in metas:
        video_views = meta["youtube_views"]

        # First filter on views
        if video_views < 100_000:
            continue

        genius_slug = meta["genius_slug"]
        for artist in artist_set:
            if fuzz.partial_ratio(artist, genius_slug) > 90:
                filtered.append(meta)
                break
    return filtered


def print_progress(start_time, processed, total, filtered):
    elapsed_time = time.time() - start_time
    estimated_total_time = (elapsed_time / processed) * total
    remaining_time = estimated_total_time - elapsed_time
    print(
        f"Processed {processed}/{total} batches. Filtered set has {filtered} items. Estimated remaining time: {remaining_time / 60:.2f} minutes",
        flush=True,
    )


if __name__ == "__main__":
    # load genius metas
    genius_metas = read_jsonl(
        "/home/christian/code/christian/metadata/genius_hq_metas.jsonl", progress=True
    )
    print("genius_metas", len(genius_metas))

    print("Loading quality map...")
    with open(
        "/home/christian/code/christian/metadata/genius_hq_metas_quality.json", "r"
    ) as f:
        metas_quality_map = json.load(f)

    print("metas_quality_map", len(metas_quality_map))

    print("Loading spot metas...")
    # spot metas contains the relevant artists to filter on
    spot_metas = read_csv("/home/christian/code/christian/metadata/spot_genres.csv")
    print("spot_metas", len(spot_metas))

    # create set of artists from spot_metas
    artist_set = set()
    for meta in spot_metas:
        artists = meta["artists"].split(";")
        artist_set.update(artists)

    print("artist_set", len(artist_set))

    # clean artist list
    artist_list = list(artist_set)
    artist_list = [artist.strip() for artist in artist_list]
    artist_list = [artist for artist in artist_list if artist != ""]
    artist_list = [artist for artist in artist_list if artist != "Various Artists"]
    # ensure length of at least 3 characters
    artist_list = [artist for artist in artist_list if len(artist) > 5]
    # ensure artist is not a number
    artist_list = [artist for artist in artist_list if not artist.isnumeric()]
    # ensure lowercase
    artist_list = [artist.lower() for artist in artist_list]

    # add "-" between spaces in artist names
    artist_list = [artist.replace(" ", "-") for artist in artist_list]

    print("cleaned artist_list", len(artist_list))

    for n in range(10):
        print(artist_list[n])

    # Use set for faster lookup
    artist_set = set(artist_list)

    # pre-filter genius metas on quality
    quality_filtered_genius_metas = []
    print()
    print("Filtering on quality...")
    for meta in tqdm(genius_metas):
        quality_meta = metas_quality_map.get(meta["id"], None)
        if quality_meta is not None:
            quality_score = float(quality_meta["audio_quality"]["score"])
            if quality_score < -0.5:
                continue
            else:
                quality_filtered_genius_metas.append(meta)

    print(f"Quality filtered metas: {len(quality_filtered_genius_metas)}")

    # Batching the genius_metas list
    batch_size = 1000  # Adjust this value as needed
    genius_metas_batches = [
        quality_filtered_genius_metas[i : i + batch_size]
        for i in range(0, len(quality_filtered_genius_metas), batch_size)
    ]

    # Create a ProcessPoolExecutor
    genius_metas_filtered = []
    start_time = time.time()

    with ProcessPoolExecutor(
        max_workers=64
    ) as executor:  # Adjust max_workers based on your CPU/core count
        futures = {
            executor.submit(filter_meta_batch, batch, artist_set): batch
            for batch in genius_metas_batches
        }

        for i, future in enumerate(as_completed(futures)):
            result = future.result()
            if result:
                genius_metas_filtered.extend(result)

            # Print progress after each future is processed
            if i % 10 == 0 and i > 0:  # Print every 10 batches
                print_progress(
                    start_time, i + 1, len(futures), len(genius_metas_filtered)
                )

    # Final progress printout
    print_progress(start_time, len(futures), len(futures), len(genius_metas_filtered))

    # The filtered list of genius metas
    print(f"Total items found: {len(genius_metas_filtered)}", flush=True)

    write_jsonl(
        genius_metas_filtered,
        "/home/christian/code/christian/metadata/genius_hq_metas_filtered.jsonl",
    )