import os
import numpy as np
from typing import Union
from concurrent.futures import ThreadPoolExecutor
import requests
import atexit

from pydantic import BaseModel, Extra

# Module-level singleton ThreadPoolExecutor
_notification_executor = ThreadPoolExecutor(max_workers=8)


# Register cleanup function to ensure ThreadPoolExecutor shuts down properly
def _cleanup_executor():
    _notification_executor.shutdown(wait=True)


atexit.register(_cleanup_executor)


class QueueItem(BaseModel, extra=Extra.forbid):
    id: str
    prompt_audio: Union[str, None] = None
    prompt_npz: Union[str, None] = None
    prompt_text: Union[str, None] = None
    metadata: dict
    # some known keys are:
    # 'type', 'stream', 'duration', 'promotion',
    # 'refund_credits',
    # ===== Experimental info =====
    # 'param_experiment': inference experiment name
    # 'model_config': inference experiment parameter values (temperature, etc)
    # 'experiment_version': logs the experiment version (v_x)
    # 'experiment": two model names in the experiment
    # ===== Prompt info =====
    # 'tags': genres
    # 'negative_tags': non-genres
    # 'prompt': lyrics
    # 'gpt_prompt': ???
    # 'gpt_description_prompt': gpt descriptions -- one box only
    # ===== Custom config =====
    # 'control_sliders': slider controls -- for custom config
    # ===== For Extra info =====
    # user_id, clip_user_id
    # ===== For audio editing =====
    # crop_start_time, crop_end_time
    # ===== Task type =====
    # task: str --
    # infill operations: "infill", "infill_intro", "infill_outro", "artist_infill", "cover_infill"
    # "artist_consistency", "cover",
    # "artist_cover"
    # "upsample"
    # extend operations: "extend", "upload_extend", "cover_extend", "artist_extend"
    # diffusion stem generation: "gen_stem"
    # diffusion stem generation: "gen_stem_2"
    # diffusion infill generation: "fixed_infill"
    # ===== For infilling =====
    # 'infill_start_s', 'infill_end_s', 'infill_context_start_s', 'infill_context_end_s'
    # 'infill_dur_s': int -- expected duration to be infilled
    # For example:
    # |   A  |   B  | I (infill) | C   | D   |
    # time at A|B: infill_context_start_s
    # time at B|I: infill_start_s
    # time at I|C: infill_end_s
    # time at C|D: infill_context_end_s
    # ===== For artist condition =====
    # 'artist_clip_id': None or str
    # 'artist_start_s': None or float
    # 'artist_end_s': None or float
    # ===== For cover condition =====
    # 'cover_clip_id': None or str
    # 'cover_start_s': None or float
    # 'cover_end_s': None or float
    # ===== For audio upload =====
    # upload_key: s3 audio file name
    # upload_type: audio upload is "audio_recording" or "file_upload"
    # is_audio_upload_tos_accepted: bool
    # has_vocal: bool
    # ===== For ditto =====
    # encode_texts_in_ditto: list of str for text encoding by ditto model
    # ===== For metrics =====
    # gen_request_start_time: int | None -- epoch millis time we received generation API request at
    gen_duration: Union[int, None] = 12
    callback_url: str | None = None
    model_name: str | None = None
    title: str | None = None

    # If this is set, this QueueItem represents a batch of items to be processed together.
    ids: list[str] | None = None

    @property
    def multi_ids(self) -> list[list[str]] | None:
        """
        multi_ids are a list of lists of ids. For example,
        [
            ["1", "2", "3"],
            ["4", "5", "6"],
        ]
        means that there are 2 requests to be generated, and each request has 3 child clips that will be generated.

        This is used for requests that generate multiple clips from a single request, such as source separation.
        """
        return self.metadata.get("multi_ids")

    @property
    def child_primary_clip_ids(self) -> list[str]:
        """The primary clip ids for the request. If a request generates multiple clips,
        this will be the first clip id from each list in multi_ids.
        """
        if self.ids:
            return self.ids
        if self.multi_ids:
            # select the first id from each list in multi_ids as the main clip id
            return [ids[0] for ids in self.multi_ids]
        return []

    @property
    def all_child_clip_ids(self) -> list[str]:
        """All ids generated by the request."""
        ids = []
        assert not (self.ids and self.multi_ids), "ids and multi_ids cannot both be set"
        if self.ids:
            ids.extend(self.ids)
        if self.multi_ids:
            for multi_id in self.multi_ids:
                if isinstance(multi_id, list):
                    ids.extend(multi_id)
                elif isinstance(multi_id, str):
                    ids.append(multi_id)
        return ids

    def notify_progress(self, result, blocking=False):
        """Send progress notification to the callback URL.

        Args:
            result: The result data to send.
            blocking: Whether to send the notification synchronously and raise exceptions.
                      Defaults to False for backward compatibility with the original behavior,
                      which didn't check for timeouts or validate response status codes.
        """
        if not self.callback_url:
            return

        headers = {
            "Authentication": f"Bearer {os.environ['API_CALLBACK_TOKEN']}",
            "X-Auth-Type": "modal",
        }

        if blocking:
            self._send_notifications_blocking(self.callback_url, result, headers)
        else:
            self._send_notifications_async(self.callback_url, result, headers)

    def _send_notifications_blocking(self, url, result, headers):
        """Send notifications synchronously, raising exceptions to caller."""
        if self.all_child_clip_ids:
            for item_id in self.all_child_clip_ids:
                requests.post(
                    url, json={"request_id": item_id, **result}, headers=headers, timeout=15
                ).raise_for_status()
        else:
            requests.post(
                url, json={"request_id": result["id"], **result}, headers=headers, timeout=15
            ).raise_for_status()

    def _send_notifications_async(self, url, result, headers):
        """Send notifications asynchronously using thread pool."""
        if self.all_child_clip_ids:
            for item_id in self.all_child_clip_ids:
                _notification_executor.submit(
                    self._send_request_async, url, {"request_id": item_id, **result}, headers
                )
        else:
            _notification_executor.submit(
                self._send_request_async, url, {"request_id": result["id"], **result}, headers
            )

    @staticmethod
    def _send_request_async(url, data, headers):
        """Fire-and-forget request with error logging."""
        try:
            requests.post(url, json=data, headers=headers, timeout=15)
        except Exception as e:
            print(f"Notification request failed: {e}")

    @property
    def is_one_box_generation(self) -> bool:
        """Check if this item is a one-box generation request."""
        gpt_description_prompt = self.metadata.get("gpt_description_prompt")
        return isinstance(gpt_description_prompt, str) and len(gpt_description_prompt) > 0

    @property
    def is_pro_generation(self) -> bool:
        """Check if this item is a pro user generation request."""
        prority = self.metadata.get("priority")
        return isinstance(prority, int) and prority > 0

    @property
    def is_free_generation(self) -> bool:
        """Check if this item is a free user generation request."""
        return not self.is_pro_generation

    @property
    def is_bot_generation(self) -> bool:
        """Check if this item is a bot generation request."""
        return self.metadata.get("is_bot", False)

    @property
    def is_image_to_song(self) -> bool:
        """Check if this item is an image to song request."""
        return (
            self.metadata.get("image_to_song_s3_ids", False)
            and len(self.metadata["image_to_song_s3_ids"]) > 0
        )

    @property
    def is_video_to_song(self) -> bool:
        """Check if this item is a video to song request."""
        return self.metadata.get("video_to_song_description") is not None

    @property
    def is_artist_condition(self) -> bool:
        """Check if this item is an artist condition generation request."""
        return (
            (
                self.metadata.get("artist_clip_id", None) is not None
                and self.metadata.get("task", "") == "artist_consistency"
            )
            or self.is_artist_infill
            or self.is_artist_extend
            or self.is_artist_cover_condition
            or self.is_artist_cover_extend
        )

    @property
    def is_cover_condition(self) -> bool:
        """Check if this item is a cover condition generation request."""
        return (
            (
                self.metadata.get("cover_clip_id", None) is not None
                and self.metadata.get("task", "") == "cover"
            )
            or self.is_cover_infill
            or self.is_cover_extend
            or self.is_artist_cover_condition
            or self.is_artist_cover_extend
        )

    @property
    def is_artist_cover_condition(self) -> bool:
        """Check if this item is an artist cover condition generation request."""
        return (
            self.metadata.get("artist_clip_id", None) is not None
            and self.metadata.get("cover_clip_id", None) is not None
            and self.metadata.get("task", "") == "artist_cover"
        )

    @property
    def is_infill(self) -> bool:
        """Check if this item is an infill generation request.

        Note that infilling curently also uses audio prompt id,
        so this is the only way telling that they are different.

        Note that artist infilling and cover infilling are also under the same check.
        """
        return (
            (
                self.metadata.get("task", "") in ["infill", "infill_intro", "infill_outro"]
                and self.prompt_audio is not None
            )
            or self.is_artist_infill
            or self.is_cover_infill
        )

    @property
    def is_upsample(self) -> bool:
        """Check if this item is an upsample generation request."""
        return self.metadata.get("task", "") == "upsample" and self.metadata.get(
            "upsample_clip_id", None
        )

    @property
    def is_stem(self) -> bool:
        """Check if this item is an diffusion stem generation request."""
        return (
            self.metadata.get("task", "") == "gen_stem" or self.metadata.get("task", "") == "gen_stem_2"
        ) and self.prompt_audio is not None

    @property
    def is_diff_seed(self) -> bool:
        """Check if this item is an diffusion seed generation request."""
        return self.metadata.get("task", "") == "seed"

    @property
    def is_diff_infill(self) -> bool:
        """Check if this item is an diffusion infill generation request."""
        # intentionally don't mention diffusion in infill in the task field
        return self.metadata.get("task", "") == "fixed_infill" and self.prompt_audio is not None

    @property
    def is_upload_extend(self) -> bool:
        """Check if this item is an upload extend generation request."""
        return self.metadata.get("task", "") == "upload_extend" and self.prompt_audio is not None

    @property
    def is_cover_extend(self) -> bool:
        """Check if this item is a cover extend generation request."""
        return (
            self.metadata.get("task", "") == "cover_extend"
            and self.prompt_audio is not None
            and self.metadata.get("cover_clip_id", None) is not None
        )

    @property
    def is_artist_extend(self) -> bool:
        """Check if this item is an artist extend generation request."""
        return (
            self.metadata.get("task", "") == "artist_extend"
            and self.prompt_audio is not None
            and self.metadata.get("artist_clip_id", None) is not None
        )

    @property
    def is_artist_cover_extend(self) -> bool:
        """Check if this item is an artist extend generation request."""
        return (
            self.metadata.get("task", "") == "artist_cover_extend"
            and self.prompt_audio is not None
            and self.metadata.get("artist_clip_id", None) is not None
            and self.metadata.get("cover_clip_id", None) is not None
        )

    @property
    def is_cover_infill(self) -> bool:
        """Check if this item is a cover infill generation request."""
        return (
            self.metadata.get("task", "") == "cover_infill"
            and self.prompt_audio is not None
            and self.metadata.get("cover_clip_id", None) is not None
        )

    @property
    def is_artist_infill(self) -> bool:
        """Check if this item is an artist infill generation request."""
        return (
            self.metadata.get("task", "") == "artist_infill"
            and self.prompt_audio is not None
            and self.metadata.get("artist_clip_id", None) is not None
        )

    @property
    def is_playlist_condition(self) -> bool:
        """Check if this item is a playlist condition generation request."""
        return (self.metadata.get("task", "") == "playlist_condition") and self.metadata.get(
            "playlist_clip_ids", None
        ) is not None

    @property
    def is_multi_artist_consistency(self) -> bool:
        """Check if this item is a multi artist consistency generation request."""
        return (self.metadata.get("task", "") == "multi_artist_consistency") and self.metadata.get(
            "artist_clip_ids", None
        ) is not None

    @property
    def is_underpainting(self) -> bool:
        """Check if this item is an underpainting generation request."""
        return (
            self.metadata.get("task", "") == "underpainting"
            and self.metadata.get("underpainting_clip_id", None) is not None
        )

    @property
    def is_overpainting(self) -> bool:
        """Check if this item is an overpainting generation request."""
        return (
            self.metadata.get("task", "") == "overpainting"
            and self.metadata.get("overpainting_clip_id", None) is not None
        )


class HistoryPrompt(BaseModel):
    prompt_audio: Union[np.ndarray, None] = None  # audio tokens
    future_audio: Union[np.ndarray, None] = None  # audio tokens of future segment
    pre_history_arr: Union[np.ndarray, None] = None  # audio before history/prompt audio
    post_future_arr: Union[np.ndarray, None] = None  # audio after prompt audio
    prompt_lyrics: Union[str, None] = None  # histoty lyrics -- just the lyrics for the history part
    cover_audio: Union[np.ndarray, None] = None  # cover audio
    artist_audio: Union[np.ndarray, None] = None  # artist audio
    underpainting_audio: Union[np.ndarray, None] = None  # underpainting audio
    overpainting_audio: Union[np.ndarray, None] = None  # overpainting audio
    playlist_audio: Union[list[np.ndarray], None] = None  # playlist audio
    multi_artist_audio: Union[list[np.ndarray], None] = None  # multi artist audio
    # history latents; for extend, infill, stem, and diffusion infilling
    history_latents: Union[np.ndarray, None] = None
    future_latents: Union[np.ndarray, None] = None  # future latents; for diffusion infilling

    class Config:
        arbitrary_types_allowed = True


class QueueItemCancel(BaseModel):
    id: str