from data import midi2wavtool, wavtool2midi
from api import ComposerAPI, GenerationRequest
from lib.redis_client import RedisClient, getenv
from math import inf
from warmup import warmup
import sys


class ComposerRedisClient(RedisClient):
    def __init__(self, model_path, *args):
        super().__init__(*args)
        self.api = ComposerAPI.from_path(model_path)

    @classmethod
    def from_env(cls):
        return super().from_env(getenv("MODEL_PATH"))

    def warmup(self):
        warmup(self.api)

    def handle_request(self, request_data, recorder, notify_progress, times_out_at):
        accompany = request_data.get("accompany", None)
        if accompany is not None:
            accompany = wavtool2midi({"notes": accompany})

        text_prompt = request_data.get("textPrompt", None)

        responses = self.api.generate(
            wavtool2midi(request_data),
            [GenerationRequest.from_json(x) for x in request_data["requests"]],
            accompany=accompany,
            text_prompt=text_prompt,
            redundant_clips=[
                wavtool2midi(x) for x in request_data.get("redundantClips", [])
            ],
            redundant_pitch_penalty=request_data.get("redundantPitchPenalty", 5.0),
            redundant_duration_penalty=request_data.get(
                "redundantDurationPenalty", 0.0
            ),
            recorder=recorder,
            deadline=(times_out_at - 0.5) if times_out_at is not None else inf,
        )

        return {"response": [midi2wavtool(x) for x in responses]}


if __name__ == "__main__":
    client = ComposerRedisClient.from_env()
    client.run()