import json
from uuid import uuid4
import redis.asyncio as aredis

import botocore.exceptions
import os
import modal
import tempfile


from suno_utils.worker.modal_base import MODAL_MOUNTS

from suno_utils.worker.modal_runner_chirp_v2 import base_image
from suno_utils.worker.loader import S3Loader
from suno_utils.worker.schema import QueueItem
from suno_utils.worker.settings import s3_client


MOUNT_PATH = "/suno/models"


SECRETS = [
    modal.Secret.from_name("studio-aws"),
    modal.Secret.from_name("redis-test"),
    modal.Secret.from_dict(
        {
            "SUNO_ASSETS_PATH": "/suno/models/assets",
            "XDG_CACHE_HOME": "/suno/models/",
        }
    ),
    modal.Secret.from_name("openai-secret"),
]


ENV_NAME = "test"
STUB_NAME = f"dummy-streaming-{ENV_NAME}"
stub = modal.Stub(STUB_NAME, image=base_image)

CHUNK_SIZE = 1024 * 64


@stub.cls(
    cpu=2.0,
    # memory=16384,
    secrets=SECRETS,
    timeout=100,
    container_idle_timeout=80,
    mounts=MODAL_MOUNTS,
    memory=4096,
    keep_warm=1,
    concurrency_limit=2,
)
class DummyStreamingStub:
    def __enter__(self):
        self.redis = aredis.Redis.from_url(os.environ.get("REDIS_TEST_URL"))
        self.loader = S3Loader()

    @modal.method()
    async def generate(self, queue_item: str):
        item = QueueItem(**json.loads(queue_item))
        audio_id = "fe8100ea-50b5-411e-bf4b-dd8d772c995b_0"

        with tempfile.NamedTemporaryFile(suffix=".mp3") as f:
            try:
                s3_client.download_fileobj("suno-data-uploads", f"studio/uploads/{audio_id}.mp3", f)
            except botocore.exceptions.ClientError:
                print("No .mp3 file found for", audio_id)

            with open(f.name, "rb") as read_f:
                data = read_f.read()

            print("Data len", len(data))

            list_id = item.id

            await self.redis.delete(list_id)

            for i in range(0, len(data), CHUNK_SIZE):
                print(i)
                await self.redis.rpush(list_id, data[i : i + CHUNK_SIZE])

            await self.redis.rpush(list_id, b"")
            await self.redis.expire(list_id, 3 * 60)

            lsize = await self.redis.llen(list_id)
            print(lsize, "chunks")

        self.loader.notify_finish(
            item,
            {
                "id": item.id,
                "model": "chirp-v2-stream-test",
                "n_audios": 1,
                "ok": 1,
                "gen_duration": 5,
                "ids": [item.id],
                "durations": [20],
            },
        )

        return item.id


@stub.local_entrypoint()
def main():
    uid = str(uuid4())
    inputs = [
        json.dumps(
            dict(
                id=uid,
                prompt_text="""Hello world""",
                metadata={"tags": "pop"},
            ),
        ),
    ]
    model = DummyStreamingStub()
    print(uid)
    for input in inputs:
        model.generate.remote(
            input,
        )