import modal
import tempfile
import time
import json

from suno_utils.worker.settings import s3_client
from suno_utils.worker.schema import QueueItem
from suno_utils.worker.modal_base import get_modal_base_image
from suno_utils.audio import Audio


############## CHANGE THESE ##############

DEPLOYMENT_TYPE = "dev"  # dev, prod
####
APP_NAME = f"whisper-{DEPLOYMENT_TYPE}"
assert APP_NAME.endswith(DEPLOYMENT_TYPE)

MODEL_DIR = "/model"
MODEL_NAME = "openai/whisper-large-v3-turbo"
MODEL_REVISION = "41f01f3fe87f28c78e2fbf8b568835947dd65ed9"

image = (
    get_modal_base_image()
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1", "HF_HUB_CACHE": MODEL_DIR})
    .pip_install(
        "torch==2.5.1",
        "transformers==4.47.1",
        "hf-transfer==0.1.8",
        "huggingface_hub==0.27.0",
        "librosa==0.10.2",
        "soundfile==0.12.1",
        "accelerate==1.2.1",
        "datasets==3.2.0",
    )
    .add_local_python_source("suno_utils", copy=False)
)

model_cache = modal.Volume.from_name("hf-hub-cache", create_if_missing=True)
app = modal.App(
    APP_NAME,
    image=image,
    volumes={MODEL_DIR: model_cache},
)


@app.function()
def download_model():
    from huggingface_hub import snapshot_download
    from transformers.utils import move_cache

    snapshot_download(
        MODEL_NAME,
        ignore_patterns=["*.pt", "*.bin"],  # Using safetensors
        revision=MODEL_REVISION,
    )
    move_cache()


SECRETS = [
    modal.Secret.from_name("studio-aws"),
    modal.Secret.from_name("redis-test"),
    modal.Secret.from_name("huggingface-secret-suno"),
    modal.Secret.from_name("api-callback-token"),
]


@app.cls(
    gpu="A10G",
    min_containers=0 if DEPLOYMENT_TYPE == "dev" else 2,
    max_containers=5 if DEPLOYMENT_TYPE == "dev" else 50,
    secrets=SECRETS,
    timeout=200,
    scaledown_window=360,
)
@modal.concurrent(max_inputs=4)
class WhisperStub:
    @modal.enter()
    def load_model(self):
        import torch
        from transformers import (
            AutoModelForSpeechSeq2Seq,
            AutoProcessor,
            pipeline,
        )

        self.processor = AutoProcessor.from_pretrained(MODEL_NAME)
        self.model = AutoModelForSpeechSeq2Seq.from_pretrained(
            MODEL_NAME,
            torch_dtype=torch.float16,
            low_cpu_mem_usage=True,
            use_safetensors=True,
        ).to("cuda")

        # Create a pipeline for preprocessing and transcribing speech data
        self.pipeline = pipeline(
            "automatic-speech-recognition",
            model=self.model,
            tokenizer=self.processor.tokenizer,
            feature_extractor=self.processor.feature_extractor,
            torch_dtype=torch.float16,
            device="cuda",
        )

    @modal.method()
    def transcribe(self, queue_item: str):
        item = QueueItem(**json.loads(queue_item))
        print(f"WhisperStub: transcribe: {queue_item}")
        s3_file_name = item.id
        start = time.monotonic_ns()
        try:
            # Create a temporary directory to store the downloaded file
            with tempfile.TemporaryDirectory() as temp_dir:
                if not s3_file_name.endswith(".mp3"):
                    s3_file_name += ".mp3"
                # Create full local path
                local_path = f"{temp_dir}/{s3_file_name}"

                # Download the file to the temporary directory
                with open(local_path, "wb") as file:
                    s3_path = "studio/uploads/" + s3_file_name
                    s3_client.download_fileobj("suno-data-uploads", s3_path, file)

                print(f"Transcribing audio samples {local_path}")

                # Load audio data from the local file path
                audio = Audio.from_file(local_path)
                audio_chunks = []
                for i in range(0, int(audio.duration_s // 30) + 1):
                    start_time = i * 30
                    end_time = (i + 1) * 30
                    if audio.duration_s < start_time:
                        break
                    audio_chunks.append(audio.get_segment(start_time, end_time))

                audio_chunk_inputs = [chunk.resample(16000).mono().array_float for chunk in audio_chunks]

                # Pass numpy array to the pipeline
                transcriptions = self.pipeline(audio_chunk_inputs)
                if not transcriptions:
                    return ""
                end = time.monotonic_ns()
                total_text = ""
                for transcription in transcriptions:
                    total_text += transcription["text"]
                print(f"Transcribed samples in {round((end - start) / 1e9, 2)}s. Text: {total_text}")
                return total_text
        except Exception as e:
            print(f"WhisperStub: transcribe: {queue_item} failed: {e}")
            return ""


@app.local_entrypoint()
def main():
    model = WhisperStub()
    model.transcribe.spawn(
        json.dumps(
            {
                "id": "ae137e1f-7562-48f6-a11c-b125e645282c",
                "metadata": {},
                "title": "Test Walking down",
            }
        )
    )
    # this is an instrumental song
    model.transcribe.spawn(
        json.dumps(
            {
                "id": "ea06bf18-d820-4b42-9b5b-54f6cf7efedb",
                "metadata": {},
                "title": "Test Instrumental",
            }
        )
    )
    print("Done")
    time.sleep(100)