from enum import Enum
from suno_utils.gpt.generation import GenerationConfig, combine_prompts
from suno_utils.gpt.generation_engine import simplify_tags, tokenize, make_audio_array
from suno_utils.gpt.modules.gpt import GPTConfig
from transformers import PreTrainedTokenizerFast
from abc import ABC
from typing import Optional, Union, List
from dataclasses import dataclass
import torch

# each prompt type inherits from GenerationPrompt, stored as PromptType
# prompt can be either a text or audio prompt, GenerationPrompt required to implement getting this from the config
# class to store all the prompts (or this could just be a list/dict??)
# class representing all the Cfg streams (CfgPromptStream)
#       a stream can be made up of one or several GenerationPrompts
#       define child classes for this that gather the full prompt tensor together (prompt, max steps and weight)
# this way we can add new prompts and new cfg streams that apply prompts

class PromptType(str, Enum):
    STYLE = "tag"
    NEG_STYLE = "neg_tag"
    LYRICS = "lyrics"
    COVER = "cover"
    ARTIST = "artist"
    FUTURE = "future"
    HISTORY = "history"
    CONTROL = "control"

@dataclass
class GenerationPrompt(ABC):
    text_prompt: Optional[str] = None
    audio_prompt: Optional[torch.Tensor] = None

    @classmethod
    def from_config(cls, gconf: GenerationConfig, model_cfg: GPTConfig, tokenizer: PreTrainedTokenizerFast):
        raise NotImplementedError

    @staticmethod
    def config_uses_prompt(gconf: GenerationConfig) -> bool:
        raise NotImplementedError

    @staticmethod
    def prompt_type() -> PromptType:
        raise NotImplementedError
    
class StylePrompt(GenerationPrompt):
    @classmethod
    def from_config(cls, gconf: GenerationConfig, model_cfg: GPTConfig, tokenizer: PreTrainedTokenizerFast):
        if not cls.config_uses_prompt(gconf):
            return cls(text_prompt="")
        
        text_tags = simplify_tags(gconf.text_tags, n_repeat=gconf.n_repeat_tags)
        return cls(text_prompt=text_tags)

    @staticmethod
    def config_uses_prompt(gconf: GenerationConfig) -> bool:
        return gconf.text_tags is not None

    @staticmethod
    def prompt_type() -> PromptType:
        return PromptType.STYLE
    
class NegStylePrompt(GenerationPrompt):
    @classmethod
    def from_config(cls, gconf: GenerationConfig, model_cfg: GPTConfig, tokenizer: PreTrainedTokenizerFast):
        if not cls.config_uses_prompt(gconf):
            return cls(text_prompt="")
        
        text_neg_tags = simplify_tags(gconf.text_neg_tags, n_repeat=gconf.n_repeat_neg_tags)
        return cls(text_prompt=text_neg_tags)

    @staticmethod
    def config_uses_prompt(gconf: GenerationConfig) -> bool:
        return gconf.text_neg_tags is not None

    @staticmethod
    def prompt_type() -> PromptType:
        return PromptType.NEG_STYLE
    

class CfgStreamPrompt(ABC):
    def __init__(self, prompt: Optional[torch.Tensor] = None):
        self.prompt = prompt

    @classmethod
    def from_config(cls, gconf: GenerationConfig, model_cfg: GPTConfig, tokenizer: PreTrainedTokenizerFast):
        if not cls.config_uses_stream(gconf):
            return cls(prompt=None)
        
        text = cls.make_text_from_config(gconf)
        tokens = tokenize(text, model_cfg, tokenizer)
        max_len = min(len(tokens), model_cfg.t_text)
        text_tensor = torch.from_numpy(tokens)[:max_len]
        audio_tensor = cls.make_audio_from_config(gconf, model_cfg)


        stream = torch.empty(
            (model_cfg.n_streams, text_tensor.shape[-1] + audio_tensor.shape[-1]), dtype=torch.long
        )
        stream[0] = model_cfg.text_pad_token
        stream[1] = model_cfg.semantic_pad_token
        stream[2:] = model_cfg.coarse_pad_token
        stream[0, : text_tensor.shape[-1]] = text_tensor

        stream[1:, text_tensor.shape[-1] :] = audio_tensor

        return cls(prompt=stream)
        
    @staticmethod
    def make_text_from_config(gconf: GenerationConfig) -> Union[str, List[str]]:
        raise NotImplementedError
    
    @staticmethod
    def make_audio_from_config(gconf: GenerationConfig, model_cfg: GPTConfig) -> torch.Tensor:
        raise NotImplementedError

    @staticmethod
    def config_uses_stream(gconf: GenerationConfig) -> bool:
        raise NotImplementedError

    @staticmethod
    def stream_type() -> CfgStreamType:
        raise NotImplementedError
    
class TextTagStreamPrompt(CfgStreamPrompt):
    @staticmethod
    def make_text_from_config(gconf: GenerationConfig) -> Union[str, List[str]]:
        text_tags = simplify_tags(gconf.text_tags, n_repeat=gconf.n_repeat_tags) if gconf.text_tags else ""
        return text_tags
    
    @staticmethod
    def make_audio_from_config(gconf: GenerationConfig, model_cfg: GPTConfig) -> torch.Tensor:
        audio_arr = make_audio_array(model_cfg, arr=gconf.history_arr)
        future_arr = make_audio_array(model_cfg, arr=gconf.future_arr, semantic_infer_token=model_cfg.semantic_future_token)

        if gconf.future_arr is not None:
            audio_arr = torch.concat([future_arr, audio_arr], dim=-1)

        if gconf.history_arr is not None:
            # crop last part
            n_history = (
                audio_arr.shape[-1]
                - model_cfg.semantic_n_codebooks * model_cfg.semantic_shift_factor
                - (model_cfg.coarse_n_codebooks - 1) * model_cfg.coarse_shift_factor
            )
            audio_arr = audio_arr[:, :n_history]
        audio_arr = audio_arr[:, : model_cfg.t_audio]

        return audio_arr

    def config_uses_stream(self, gconf: GenerationConfig) -> bool:
        return gconf.text_tags is not None and gconf.cfg_coef_tags != 0

    def stream_type(self) -> CfgStreamType:
        return CfgStreamType.TAG
    
class NegTextTagStreamPrompt(CfgStreamPrompt):
    @staticmethod
    def make_text_from_config(gconf: GenerationConfig) -> Union[str, List[str]]:
        text_neg_tags = (
            simplify_tags(gconf.text_neg_tags, n_repeat=gconf.n_repeat_neg_tags)
            if gconf.text_neg_tags
            else ""
        )
        return text_neg_tags
    
    @staticmethod
    def make_audio_from_config(gconf: GenerationConfig, model_cfg: GPTConfig) -> torch.Tensor:
        audio_arr = make_audio_array(model_cfg, arr=gconf.history_arr)
        future_arr = make_audio_array(model_cfg, arr=gconf.future_arr, semantic_infer_token=model_cfg.semantic_future_token)

        if gconf.future_arr is not None:
            audio_arr = torch.concat([future_arr, audio_arr], dim=-1)

        if gconf.history_arr is not None:
            # crop last part
            n_history = (
                audio_arr.shape[-1]
                - model_cfg.semantic_n_codebooks * model_cfg.semantic_shift_factor
                - (model_cfg.coarse_n_codebooks - 1) * model_cfg.coarse_shift_factor
            )
            audio_arr = audio_arr[:, :n_history]
        audio_arr = audio_arr[:, : model_cfg.t_audio]

        return audio_arr

    @staticmethod
    def config_uses_stream(gconf: GenerationConfig) -> bool:
        return gconf.text_neg_tags is not None and gconf.cfg_coef_neg_tags != 0

    @staticmethod
    def stream_type() -> CfgStreamType:
        return CfgStreamType.NEG_TAG
    
class MainStreamPrompt(CfgStreamPrompt):
    @classmethod
    def from_config(cls, gconf: GenerationConfig):
        # Prepare main text
        text = gconf.text.strip() if gconf.text else ""
        text = text.replace("\n", " ") if not gconf.respect_newlines else text
        text = f"{gconf.history_text.strip()}\n\n{text}" if gconf.history_text else text

        # Process tags
        text_tags = simplify_tags(gconf.text_tags, n_repeat=gconf.n_repeat_tags) if gconf.text_tags else ""

        # Combine text with tags if applicable
        if text_tags:
            text = combine_prompts(text_tags, text)

        # Combine prompts with control tags
        if gconf.text_start_control_tags:
            text_control_tag = gconf.text_start_control_tags
            text = combine_prompts(text_control_tag, text)

        if gconf.text_end_control_tags:
            text_control_tag = gconf.text_end_control_tags
            text = combine_prompts(text, text_control_tag)

    def stream_type(self):
        return CfgStreamType.MAIN
    
class CfgPrompt:
    def __init__(self, gen_config: GenerationConfig, stream_type: CfgStreamType):
        self.stream_type = type

    @staticmethod
    def make_text_prompt(gconf: GenerationConfig, tokenizer: PreTrainedTokenizerFast, stream_type: CfgStreamType):