"""Retrieve existing lyrics by id"""

import logging
import os
import time
import json

import modal
from suno_utils.utils.s3 import _download_s3_file
from suno_utils.worker.loader import S3Loader
from suno_utils.worker.modal_base import MODAL_MOUNTS, get_modal_base_image
from suno_utils.worker.settings import s3_client
from suno_utils.worker.utils import retry_decorator

logger = logging.getLogger(__name__)
logging.basicConfig()
logger.setLevel(logging.INFO)

############## CHANGE THESE ##############
DEPLOYMENT_TYPE = "dev"

##########################################
DEPLOYMENT_TYPES = {"dev", "prod"}
assert DEPLOYMENT_TYPE in DEPLOYMENT_TYPES
APP_NAME = f"retrieve-existing-lyrics-by-id-{DEPLOYMENT_TYPE}"
assert APP_NAME.endswith(DEPLOYMENT_TYPE)
MOUNT_PATH = "/suno/models"
UPLOADS_S3_BUCKET = "suno-data-uploads"
LYRICS_FILE_NAME = "noempty_documents"
LYRICS_FILE_PATH = os.path.join(MOUNT_PATH, f"{LYRICS_FILE_NAME}.jsonl")
POPULAR_LYRICS_FILE_NAME = "popular_lyrics"
POPULAR_LYRICS_FILE_PATH = os.path.join(MOUNT_PATH, f"{POPULAR_LYRICS_FILE_NAME}.jsonl")

retry_s3_download = retry_decorator(3, wait_seconds=20)(s3_client.download_fileobj)

N_CPU = 1
aws_secret = modal.Secret.from_name("studio-aws")
DD_SAMPLE_RATE = "1" if DEPLOYMENT_TYPE == "dev" else "0.0001"
SECRETS = [
    aws_secret,
    modal.Secret.from_dict(
        {
            "SUNO_ASSETS_PATH": "/suno/models/assets",
            "XDG_CACHE_HOME": "/suno/models/",
        },
    ),
    modal.Secret.from_name("openai-secret"),
    modal.Secret.from_dict(
        {
            "DD_SITE": "datadoghq.com",
            "DD_ENV": DEPLOYMENT_TYPE,
            "DD_SERVICE": "chatgpt-worker",
            "DD_LOGS_ENABLED": "true",
            "DD_TRACE_ENABLED": "true",
            "DD_TRACE_SAMPLE_RATE": DD_SAMPLE_RATE,
        },
    ),
    modal.Secret.from_name("datadog-metrics"),
    modal.Secret.from_name("api-callback-token"),
]


class LyricsRetrieverWorker(S3Loader):
    """Retrieve existing lyrics by id"""

    def __init__(self, file_path: str):
        self._file_path = file_path
        self._documents = []
        self._popular_documents = []

    def preload(self):
        print(f"Preloading documents from {self._file_path}")
        start_time = time.time()
        with open(self._file_path, "r") as json_file:
            self._documents = json.load(json_file)
        finish_time = time.time()
        print(f"Preloaded {len(self._documents)} documents in {finish_time - start_time}s")
        start_time = finish_time
        print(f"Preloading popular documents from {self._file_path}")
        with open(self._file_path, "r") as json_file:
            self._popular_documents = json.load(json_file)
        finish_time = time.time()
        print(
            f"Preloaded {len(self._popular_documents)} popular documents in {finish_time - start_time}s"
        )

    def get_lyrics_by_id(self, lyric_id):
        if lyric_id.startswith("pdoc_"):
            id = int(lyric_id.split("_")[-1])
            if id >= len(self._popular_documents):
                return None
            return self._popular_documents[id]
        else:
            id = int(lyric_id.split("_")[-1])
            if id >= len(self._documents):
                return None
            return self._documents[id]


def download_document_wrapper() -> None:
    start_time = time.time()
    target_lyrics_path = os.path.join(MOUNT_PATH, f"{LYRICS_FILE_NAME}.jsonl")
    print("Downloading lyrics files")
    _download_s3_file(
        f"s3://suno-data/ashe/trained_models/{LYRICS_FILE_NAME}.jsonl",
        target_lyrics_path,
    )
    finish_time = time.time()
    print(f"Downloaded lyrics files successfully in {finish_time - start_time}s")
    start_time = finish_time
    print("Downloading popular lyrics files")
    _download_s3_file(
        f"s3://suno-data/ashe/trained_models/{POPULAR_LYRICS_FILE_NAME}.jsonl",
        target_lyrics_path,
    )
    finish_time = time.time()
    print(f"Downloaded popular lyrics files successfully in {finish_time - start_time}s")


base_image = get_modal_base_image()
image = base_image.run_function(download_document_wrapper, secrets=SECRETS)
app = modal.App(APP_NAME, image=image)


@app.cls(
    cpu=N_CPU,
    secrets=SECRETS,
    timeout=1000,
    scaledown_window=360,
    mounts=MODAL_MOUNTS,
    retries=modal.Retries(
        max_retries=2,
        backoff_coefficient=2.0,
        initial_delay=5.0,
    ),
    allow_concurrent_inputs=20,
    min_containers=1,
)
class LyricsRetrieverApp:
    """App for retrieving existing lyrics by id"""

    def __init__(self):
        """Set up LyricsRetrieverApp."""
        self.worker = LyricsRetrieverWorker(LYRICS_FILE_PATH)
        self.worker.preload()

    @modal.method()
    def retrieve_existing_lyrics_by_id(self, lyrics_id: str):
        return self.worker.get_lyrics_by_id(lyrics_id)


@app.local_entrypoint()
def main():
    testApp = LyricsRetrieverApp()
    test_id = "doc_123"
    result = testApp.retrieve_existing_lyrics_by_id.remote(test_id)
    print(f"lyrics copyright infringement result: {result}")

    test_id = "pdoc_256328"
    result = testApp.retrieve_existing_lyrics_by_id.remote(test_id)
    print(f"lyrics copyright infringement result: {result}")