import modal
from suno_utils.audio import Audio
from suno_utils.worker.audio_fingerprints.shazam import AudioFeatureExtractor
from suno_utils.worker.modal_base import get_modal_base_image
import platform
import boto3
import pickle
import numpy as np
import logging
from typing import Optional
import concurrent.futures
from collections import Counter
import time
import asyncio

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

DEPLOYMENT_TYPE = "dev"  # dev, prod

##########################################
image = get_modal_base_image()
SECRETS = [
    modal.Secret.from_name("studio-aws"),
]
TOTAL_PARTITIONS = 160

APP_NAME = f"audio-search-orchestrator-{DEPLOYMENT_TYPE}"
platform.node = lambda: "suno-patched-platform-node"
app = modal.App(APP_NAME, image=image)


@app.cls(
    image=image,
    secrets=SECRETS,
    cpu=4,
    memory=60000,
    min_containers=1,
)
class AudioFeatureOchestrator:
    """
    This class is used to orchestrate the audio feature extraction process.

    It will spawn a number of AudioFingerprintApp instances, and distribute the work among them.
    For every query, it will distribute the hash lookups to the different instances, and then merge the results.
    """

    def __init__(self):
        self.index_apps = [None for _ in range(TOTAL_PARTITIONS)]
        self.index_app_cls = modal.Cls.from_name(
            f"genius-hash-index-app-{DEPLOYMENT_TYPE}", "AudioFeatureIndexApp"
        )
        for partition in range(TOTAL_PARTITIONS):
            # Check if the index app for this partition already exists
            try:
                app_stats = self.index_app_cls(partition=partition).query.get_current_stats()
                print(f"Index app for partition {partition} already exists, stats: {app_stats}")
                if app_stats.num_total_runners == 0:
                    self.index_app_cls(partition=partition).query.spawn(hash=100)
                self.index_apps[partition] = self.index_app_cls(partition=partition)
            except Exception as e:
                print(
                    f"Index app for partition {partition} does not exist, exception: {e}, creating new one..."
                )
                self.index_app_cls(partition=partition).query.spawn(hash=100)
                self.index_apps[partition] = self.index_app_cls(partition=partition)
        self.audio_feature_extractor = AudioFeatureExtractor(device="cpu", fan_value=5, kernal_size=5)

    @modal.method()
    def query(
        self,
        audio_waveform: Optional[np.ndarray] = None,
        num_hashes: int = 500,
        threshold: int = 30,
        timeout: float = 10.0,
    ):
        hash_features = self.audio_feature_extractor.fingerprint(audio_waveform, delta_compress=False)
        hash_features = hash_features[:num_hashes]
        print(f"Hash features: total length: {len(hash_features)}, Example: {hash_features[:5]}")
        start = time.time()
        candidates = []

        with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
            # Submit all tasks
            future_to_partition = {
                executor.submit(
                    lambda p: self.index_apps[p]
                    .query_many.spawn(hashes=hash_features, threshold=threshold)
                    .get(),
                    partition_id,
                ): partition_id
                for partition_id in range(TOTAL_PARTITIONS)
            }

            # Wait for completion with timeout
            done, not_done = concurrent.futures.wait(
                future_to_partition.keys(), timeout=timeout, return_when=concurrent.futures.ALL_COMPLETED
            )

            # Collect completed results
            for future in done:
                try:
                    result = future.result()
                    if result is not None:
                        candidates.extend(result)
                        partition_id = future_to_partition[future]
                        print(f"Result {partition_id}: {result}")
                except Exception as e:
                    print(f"An error occurred: {e}")

            # Cancel pending futures
            for future in not_done:
                future.cancel()
            # Soer the candidates by count in descending order
            candidates = sorted(candidates, key=lambda x: x[1], reverse=True)
            print(
                f"Completed: {len(done)}, Timed out: {len(not_done)}, time taken: {time.time() - start}"
            )
            return candidates

    @modal.method()
    def get_audio_from_s3(self, s3_path: str):
        return Audio.from_s3(s3_path, sample_rate=16000)


@app.local_entrypoint()
def main():
    # test the audio feature extractor
    orchestrator = AudioFeatureOchestrator()
    audio_waveform = orchestrator.get_audio_from_s3.remote(
        "s3://suno-data-uploads/studio/uploads/b15707c2-c0e8-45f0-a0b4-a32db517fe1a.mp3"
    )
    result = orchestrator.query.remote(audio_waveform=audio_waveform.array_float)
    print(result)