import modal
import logging
from suno_utils.worker.modal_base import get_modal_base_image
import platform
import boto3
import pickle

import tempfile
from collections import Counter
import time
import numpy as np


logger = logging.getLogger(__name__)
############## CHANGE THESE ##############

DEPLOYMENT_TYPE = "dev"  # dev, prod

##########################################
image = get_modal_base_image().apt_install("build-essential")

SECRETS = [
    modal.Secret.from_name("studio-aws"),
]

fingerprint_volume = modal.Volume.from_name("audio-fingerprint-volume", create_if_missing=True)
APP_NAME = f"genius-hash-index-app-{DEPLOYMENT_TYPE}"
platform.node = lambda: "suno-patched-platform-node"
app = modal.App(APP_NAME, image=image)


def read_pickle_from_s3(bucket: str, path: str) -> dict:
    with tempfile.NamedTemporaryFile(suffix=".pkl") as f:
        start = time.time()
        S3_CLIENT = boto3.client("s3")
        S3_CLIENT.download_fileobj(Bucket=bucket, Key=path, Fileobj=f)
        print(f"Time taken to download pickle from s3: {time.time() - start}")
        start = time.time()
        # IMPORTANT: Seek to beginning of file before reading
        f.seek(0)
        data_obj = pickle.load(f)
        print(f"Time taken to read pickle: {time.time() - start}")
    return data_obj


def read_pickle_from_volume(path: str) -> dict:
    t0 = time.perf_counter()
    with open(path, "rb", buffering=1 << 20) as f:
        raw = f.read()
    t1 = time.perf_counter()
    obj = pickle.loads(raw)
    t2 = time.perf_counter()
    print(
        f"Time taken to read pickle from volume: {t2 - t0}, time to read raw: {t1 - t0}, time to load: {t2 - t1}"
    )
    return obj


@app.cls(
    image=image,
    secrets=SECRETS,
    cpu=12,
    memory=80000,
    volumes={"/audio-fingerprint-volume": fingerprint_volume},
    max_containers=200,
    timeout=3000,
)
@modal.concurrent(max_inputs=10)
class AudioFeatureIndexApp:
    """
    This class is used to query the genius hash index

    In the index, each key is a hash and the value is a list of tuples, where each tuple contains a song id (uuid) and a time offset in integer.
    The index is partitioned by shard, and each shard is stored in a separate file in s3.
    """

    partition: int = modal.parameter()

    @modal.enter()
    def load_index(self):
        """
        Initialize the index
        """
        self._index = {}
        # Load the index from s3: s3://suno-data/ashe/processed_genius_fingerprints/genius_shazam_index_dict_shard_{partition}.pkl
        print(
            f"Loading partition {self.partition} index from s3 file: ashe/genius_fingerprints/genius_shazam_{self.partition}_dict.pkl... "
        )
        # self._index = read_pickle_from_s3("suno-data",f"ashe/genius_fingerprints/genius_shazam_{self.partition}_dict.pkl")
        self._index = read_pickle_from_volume(
            f"/audio-fingerprint-volume/partition_genius_shazam_{self.partition}_dict.pkl"
        )
        print(f"Index loaded partition {self.partition} from s3 ... ")

        start = time.time()
        self._rebuild_2d_index()
        print(f"Rebuilt 2d index, took {time.time() - start} seconds")

        # Dump the 2d index to volumn too
        with open(
            f"/audio-fingerprint-volume/partition_genius_shazam_{self.partition}_2d_index.pkl", "wb"
        ) as f:
            pickle.dump(self._2d_index, f)
        fingerprint_volume.commit()

        # Optional, build a counter index
        print("Building counter index...")
        start = time.time()
        self._counter_index = {}
        for key in self._index:
            self._counter_index[key] = Counter([song_id for song_id, _ in self._index[key]])
        print(f"Done building counter index, took {time.time() - start} seconds")

    def _prune(self):
        """
        Prune the index
        """
        sizes = [len(self._index[key]) for key in self._index]
        SIZE_THRESHOLD = np.percentile(sizes, 99)
        print(f"P99 length of index is {SIZE_THRESHOLD}")
        for key in self._index:
            if len(self._index[key]) >= SIZE_THRESHOLD:
                self._index[key] = []

    def _rebuild_2d_index(self):
        """
        Rebuild the 2d index
        """
        self._2d_index = {}
        for key in self._index:
            self._2d_index[key] = {}
            for song_id, offset_list in self._index[key]:
                self._2d_index[key][song_id] = offset_list

    @modal.method()
    def query(self, hash: int) -> list[tuple[str, int]]:
        return self._index.get(hash, [])

    @modal.method()
    def query_many(self, hashes: list[tuple[int, int]], threshold: int = 1):
        """
        Query the index for multiple hashes and return a Counter of the results grouped by (song_id, offset)
        """
        print(
            f"Partition {self.partition} querying {len(hashes)} hashes, index size: {len(self._index)}"
        )
        start = time.time()
        # Step 1 find the most common song_ids
        most_common_song_ids = Counter()
        for hash, _ in hashes:
            if hash in self._counter_index:
                most_common_song_ids.update(self._counter_index[hash])
        # candidate song_ids are song_ids occur more than threshold times
        candidate_song_ids = [
            song_id for song_id, count in most_common_song_ids.items() if count > len(hashes) * 0.2
        ]
        print(
            f"Took {time.time() - start} to find candidate song_ids {len(candidate_song_ids)} with threshold {len(hashes) * 0.2}"
        )
        start = time.time()

        results = Counter()
        results = Counter(
            [
                (song_id, offset - t0)
                for h, t0 in hashes
                for song_id in candidate_song_ids
                for offset in self._2d_index.get(h, {}).get(song_id, [])
            ]
        )
        print(
            f"Time taken to query {len(hashes)} hashes: {time.time() - start}, length of results: {len(results)}"
        )
        start = time.time()
        result = Counter({k: v for k, v in results.items() if v > threshold}).most_common(5)
        print(
            f"Time taken to filter results: {time.time() - start}, length of filtered results: {len(result)}"
        )
        return result

    @modal.method()
    def query_many_fast(self, hashes, threshold=1):
        from suno_utils.worker.audio_fingerprints.fast_query import query_many

        return query_many(hashes, self._index, threshold)


@app.function(volumes={"/audio-fingerprint-volume": fingerprint_volume}, secrets=SECRETS)
def download_indexfile_from_s3(partition: int, bucket: str, key: str, file_name: str):
    """
    Download the index file from s3
    """
    import os

    S3_CLIENT = boto3.client("s3")
    os.makedirs("/audio-fingerprint-volume", exist_ok=True)
    print(f"Downloading index file for partition {partition} from s3 ... ")
    with open(f"/audio-fingerprint-volume/{file_name}", "wb") as f:
        S3_CLIENT.download_fileobj(Bucket=bucket, Key=key, Fileobj=f)
        f.seek(0)
    fingerprint_volume.commit()
    print(f"Downloaded index file for partition {partition} from s3 ... ")


@app.local_entrypoint()
def main():
    # # download all the index files from s3 to the volume
    # for partition in range(160):
    #     download_indexfile_from_s3.remote(partition, "suno-data", f"ashe/genius_fingerprints/genius_shazam_{partition}_dict.pkl", f"partition_genius_shazam_{partition}_dict.pkl")

    hashes = [
        (186580993, 1),
        (185532418, 1),
        (184680451, 1),
        (523436032, 2),
        (521076737, 2),
        (520224770, 2),
        (520683523, 2),
        (251789313, 3),
        (252248066, 3),
        (252903426, 3),
        (253165570, 3),
        (34144257, 4),
        (34799617, 4),
        (35061761, 4),
        (34471938, 4),
        (152240128, 5),
        (152502272, 5),
        (151912449, 5),
        (153485313, 5),
        (320274432, 5),
        (319684609, 5),
        (321257473, 5),
        (320602114, 5),
        (386793473, 5),
        (388366337, 5),
        (387710978, 5),
        (386990085, 5),
        (237371392, 6),
        (236716033, 6),
        (235995140, 6),
        (236912644, 6),
        (639369217, 6),
        (638648324, 6),
        (639565828, 6),
        (639827972, 6),
        (470876163, 7),
        (471793667, 7),
        (472055811, 7),
        (471269380, 7),
        (287244288, 10),
        (287506432, 10),
        (286720001, 10),
        (287703041, 10),
        (522387456, 10),
        (521601025, 10),
        (522584065, 10),
        (520224771, 10),
        (588709889, 10),
        (589692929, 10),
        (588644356, 10),
        (388366336, 11),
        (386007042, 11),
        (387317763, 11),
        (388235268, 11),
        (638976003, 11),
        (639893508, 11),
        (640876550, 11),
        (34996225, 13),
        (33751045, 13),
        (371458049, 14),
        (372441091, 14),
        (369295364, 14),
        (370409476, 14),
        (607322114, 15),
        (605290499, 15),
        (606142467, 15),
        (857800705, 17),
        (51642368, 18),
        (52494336, 18),
        (50331649, 18),
        (337707008, 18),
        (335544321, 18),
        (335740931, 18),
        (553844739, 18),
        (196610, 19),
        (4006871041, 20),
        (4007329793, 20),
        (52363267, 21),
        (3571122176, 21),
        (522452992, 24),
        (521273345, 24),
        (521601025, 24),
        (523829249, 24),
        (605159425, 24),
        (605487105, 24),
        (607715329, 24),
        (303497216, 25),
        (961085440, 25),
        (1229062145, 25),
        (1869742080, 25),
        (1110114305, 26),
        (1111228417, 26),
        (1113522177, 26),
        (725352448, 27),
        (723779585, 27),
        (1008992257, 27),
        (1010040833, 27),
        (1600454656, 27),
        (1598488577, 27),
        (1699151873, 27),
        (607387648, 28),
        (605683713, 28),
        (877068288, 28),
        (874119169, 28),
        (874577921, 28),
        (1194393601, 28),
        (52035584, 29),
        (52494336, 29),
        (438370304, 29),
        (439418880, 29),
        (556859392, 29),
        (827260928, 29),
        (1331298304, 29),
        (1329266689, 29),
        (1517682688, 29),
        (1514799105, 29),
        (994705408, 30),
        (993263617, 30),
        (994115585, 30),
        (1244921857, 30),
        (1245773825, 30),
        (876675072, 31),
        (873922561, 31),
        (874250241, 31),
        (1093074945, 31),
        (51838976, 32),
        (52166656, 32),
        (387710976, 32),
        (388431872, 32),
        (387383299, 32),
        (472317952, 32),
        (473497602, 32),
        (471269379, 32),
        (469958660, 32),
        (658046978, 32),
        (655818755, 32),
        (654966791, 32),
        (386072577, 35),
        (386531332, 35),
        (386072581, 35),
        (389218311, 35),
        (50987011, 36),
        (50528260, 36),
        (50528263, 36),
        (167968769, 39),
        (167968772, 39),
        (50528259, 40),
        (50528260, 43),
        (50528264, 43),
        (3705733120, 45),
        (50528260, 47),
        (51052550, 47),
        (50528263, 47),
        (51380231, 47),
        (51052546, 51),
        (50528259, 51),
        (51380227, 51),
        (184745985, 53),
        (185597953, 53),
        (186974210, 53),
        (185925635, 53),
        (51380224, 54),
        (51707906, 54),
        (51052547, 54),
        (270860289, 54),
        (269811714, 54),
        (269156355, 54),
        (268632068, 54),
        (622133249, 55),
        (621477890, 55),
        (623378435, 55),
        (353042433, 56),
        (352518146, 56),
        (354942978, 56),
        (353632260, 56),
        (184745985, 57),
        (187170817, 57),
        (185860099, 57),
        (184745988, 57),
        (51642370, 58),
        (50528259, 58),
        (50528263, 58),
        (672399362, 58),
        (672595975, 58),
        (335740929, 60),
        (335740933, 60),
        (337051653, 60),
        (50528260, 61),
        (51838980, 61),
        (51052550, 61),
        (51838976, 65),
        (51052546, 65),
        (389218304, 65),
        (386596866, 65),
        (389218307, 65),
        (386072580, 65),
        (858980355, 65),
        (184745986, 67),
        (185597954, 67),
        (184745989, 67),
        (51380224, 69),
        (50528259, 69),
        (51576835, 69),
        (268632067, 69),
        (269680643, 69),
        (270860292, 69),
        (271056902, 69),
        (51576832, 72),
        (50528260, 72),
        (321191937, 72),
        (321388547, 72),
        (318963716, 72),
        (320536580, 72),
        (623378434, 73),
        (622526467, 73),
        (622723075, 73),
        (672858113, 75),
        (673054721, 75),
        (672268290, 75),
        (52101120, 76),
        (52297728, 76),
        (51511297, 76),
        (51838978, 76),
        (454950912, 76),
        (454164481, 76),
        (454492162, 76),
        (454819843, 76),
        (504496129, 76),
        (504823810, 76),
        (505151491, 76),
        (504561669, 76),
        (303497217, 77),
        (303824898, 77),
        (303235076, 77),
        (304480261, 77),
        (387710977, 78),
        (387121155, 78),
        (388366340, 78),
        (388759556, 78),
        (471007234, 79),
        (472252419, 79),
        (472645635, 79),
        (472907780, 79),
        (321257473, 81),
        (321650689, 81),
        (321912834, 81),
        (640417792, 82),
        (640679937, 82),
        (641335297, 82),
        (641597441, 82),
        (741343233, 82),
        (741998593, 82),
        (742260737, 82),
        (742916097, 82),
        (809107456, 83),
        (809369600, 83),
        (810024960, 83),
        (810287104, 83),
        (977141760, 83),
        (977797120, 83),
        (978059264, 83),
        (978583552, 83),
        (1044905984, 83),
        (1045168128, 83),
        (1045692416, 83),
        (1212940288, 83),
        (1213464576, 83),
        (1280573440, 83),
        (521404418, 84),
        (520683525, 84),
        (521666565, 84),
        (520224774, 84),
        (336134147, 86),
        (337117187, 86),
        (335675396, 86),
        (337903620, 86),
        (152567808, 89),
        (151126017, 89),
        (153354241, 89),
        (402784257, 89),
        (405012481, 89),
        (405274626, 89),
        (404029444, 89),
        (34930691, 90),
        (34603012, 90),
        (606601217, 90),
        (605356035, 90),
        (605028356, 90),
        (605618181, 90),
        (672464898, 91),
        (672137219, 91),
        (672727044, 91),
        (673054725, 91),
        (353370113, 93),
        (353959938, 93),
        (354287619, 93),
        (355139587, 93),
        (270073857, 94),
        (270401538, 94),
        (271253506, 94),
        (421396481, 95),
        (422248449, 95),
        (421724162, 95),
        (506134528, 96),
        (505610241, 96),
        (503447554, 96),
        (725483520, 96),
        (723714049, 96),
        (723189764, 96),
        (1042481153, 96),
        (1043005444, 96),
        (588972035, 97),
        (590020611, 97),
        (35323906, 98),
        (33751044, 98),
        (455802880, 100),
        (453181442, 100),
        (455802883, 100),
        (725483522, 100),
        (724238339, 100),
        (50331650, 102),
        (50528260, 102),
        (1043005441, 102),
        (723320835, 103),
        (725483525, 103),
        (196610, 104),
        (1900546, 104),
        (52232192, 106),
        (50528260, 106),
        (489357315, 106),
        (486735876, 106),
        (487194629, 106),
        (1043005441, 108),
        (722403332, 109),
        (723320836, 109),
        (50987009, 110),
        (51314691, 110),
        (52232195, 110),
        (168755202, 111),
        (169672706, 111),
        (167968771, 111),
        (253558784, 113),
        (254476288, 113),
        (251854849, 113),
        (253362180, 113),
        (489357312, 113),
        (486735873, 113),
        (488243204, 113),
        (486735877, 113),
        (723124228, 113),
        (724238341, 113),
        (52035587, 114),
        (50528260, 114),
        (51183621, 114),
        (436404225, 117),
        (439025665, 117),
        (437059586, 117),
        (436404228, 117),
        (51183617, 118),
        (50528259, 118),
        (52297732, 118),
        (722272257, 118),
        (723386372, 118),
        (218300418, 119),
        (220069891, 119),
        (218300422, 119),
        (220921863, 119),
        (52297729, 121),
        (50528260, 121),
        (51314694, 121),
        (503513091, 122),
        (506134532, 122),
        (504299525, 122),
        (503513094, 122),
        (51314690, 125),
        (50528259, 125),
        (722403329, 126),
        (724238341, 126),
        (251854849, 127),
        (254476292, 127),
        (251854853, 127),
        (254476295, 127),
        (50528260, 128),
        (50528264, 128),
        (724238339, 131),
        (724107271, 131),
        (50528260, 132),
        (724107268, 134),
        (725942276, 134),
        (724697093, 134),
        (692387840, 138),
        (691142657, 138),
        (691666945, 138),
        (689307650, 138),
        (1160904705, 138),
        (1161428993, 138),
        (842661888, 139),
        (840302593, 139),
        (842858498, 139),
        (977076226, 139),
        (370016257, 140),
        (370278402, 140),
        (369295363, 140),
        (236060673, 141),
        (235077634, 141),
        (1028718594, 141),
        (1028915202, 141),
        (302186497, 142),
        (304349187, 142),
        (1364459520, 143),
        (606601216, 145),
        (606928896, 145),
        (607256576, 145),
        (605421569, 145),
        (674037760, 145),
        (674365440, 145),
        (672530433, 145),
        (758251520, 145),
        (756416513, 145),
        (756809731, 145),
        (840302593, 145),
        (840695811, 145),
        (369295361, 146),
        (370933762, 146),
        (369295364, 146),
        (370081797, 146),
        (52166657, 147),
        (50528259, 147),
        (51314692, 147),
        (52166662, 147),
        (469958658, 148),
        (470745091, 148),
        (471597061, 148),
        (469958662, 148),
        (51314689, 150),
        (52166659, 150),
        (50528260, 150),
        (51773444, 150),
        (253493250, 151),
        (251854851, 151),
        (253100035, 151),
        (252379140, 151),
        (469958657, 153),
        (471203841, 153),
        (470482946, 153),
        (469958660, 153),
        (51773440, 154),
        (51052545, 154),
        (50528259, 154),
        (369819649, 154),
        (369295363, 154),
        (372375556, 154),
        (369295367, 154),
        (184745986, 155),
        (184745990, 155),
        (185991174, 155),
        (50528260, 157),
        (51773444, 157),
        (50528264, 157),
        (840302595, 158),
        (51773440, 161),
        (50528260, 161),
        (50528263, 161),
        (51773448, 161),
        (369295364, 161),
        (369295367, 161),
        (370540552, 161),
        (370933769, 161),
        (50528259, 165),
        (51773444, 165),
        (52166661, 165),
        (50987014, 165),
        (51773441, 168),
        (52166658, 168),
        (50987011, 168),
        (51314691, 168),
        (370933761, 169),
        (369754114, 169),
        (370081794, 169),
        (372375554, 169),
        (470417409, 170),
        (470745089, 170),
        (473038849, 170),
        (469958658, 170),
        (168755200, 171),
        (167968769, 171),
        (169148417, 171),
        (251854849, 171),
        (253034497, 171),
        (253952002, 171),
        (840237057, 171),
        (841154562, 171),
        (840695812, 171),
        (51707904, 172),
        (52166659, 172),
        (50528260, 172),
        (354615297, 172),
        (354156547, 172),
        (352518148, 172),
        (353763333, 172),
    ]
    index_app = AudioFeatureIndexApp(partition=0)
    result = index_app.query_many.remote(hashes=hashes, threshold=50)
    print(f"Result using query_many: {result}")