import time
import uuid
from enum import Enum
from typing import Any, ClassVar, Dict, List, Optional, OrderedDict, Set, Union

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

from predibase.resources.lorax.errors import ValidationError

ADAPTER_SOURCES = ["hub", "local", "s3", "pbase"]
MERGE_STRATEGIES = ["linear", "ties", "dare_linear", "dare_ties"]
MAJORITY_SIGN_METHODS = ["total", "frequency"]


class MergedAdapters(BaseModel):
    # IDs of the adapters to merge
    ids: List[str]
    # Weights of the adapters to merge
    weights: List[float]
    # Merge strategy
    merge_strategy: Optional[str] = None
    # Density
    density: float
    # Majority sign method
    majority_sign_method: Optional[str] = None

    @field_validator("ids")
    @classmethod
    def validate_ids(cls, v):
        if not v:
            raise ValidationError("`ids` cannot be empty")
        return v

    @field_validator("weights")
    @classmethod
    def validate_weights(cls, v, values):
        ids = values.data["ids"]
        if not v:
            raise ValidationError("`weights` cannot be empty")
        if len(ids) != len(v):
            raise ValidationError("`ids` and `weights` must have the same length")
        return v

    @field_validator("merge_strategy")
    @classmethod
    def validate_merge_strategy(cls, v):
        if v is not None and v not in MERGE_STRATEGIES:
            raise ValidationError(f"`merge_strategy` must be one of {MERGE_STRATEGIES}")
        return v

    @field_validator("density")
    @classmethod
    def validate_density(cls, v):
        if v < 0 or v > 1.0:
            raise ValidationError("`density` must be >= 0.0 and <= 1.0")
        return v

    @field_validator("majority_sign_method")
    @classmethod
    def validate_majority_sign_method(cls, v):
        if v is not None and v not in MAJORITY_SIGN_METHODS:
            raise ValidationError(
                f"`majority_sign_method` must be one of {MAJORITY_SIGN_METHODS}",
            )
        return v


class ResponseFormatType(str, Enum):
    json_object = "json_object"


class ResponseFormat(BaseModel):
    model_config = ConfigDict(use_enum_values=True)

    type: ResponseFormatType
    schema_spec: Optional[Union[Dict[str, Any], OrderedDict]] = Field(
        None,
        alias="schema",
    )


class Parameters(BaseModel):
    # The ID of the adapter to use
    adapter_id: Optional[str] = None
    # The source of the adapter to use
    adapter_source: Optional[str] = None
    # Adapter merge parameters
    merged_adapters: Optional[MergedAdapters] = None
    # API token for accessing private adapters
    api_token: Optional[str] = None
    # Activate logits sampling
    do_sample: bool = False
    # Maximum number of generated tokens
    max_new_tokens: Optional[int] = None
    # Whether to ignore the EOS token during generation
    ignore_eos_token: bool = False
    # The parameter for repetition penalty. 1.0 means no penalty.
    # See [this paper](https://arxiv.org/pdf/1909.05858.pdf) for more details.
    repetition_penalty: Optional[float] = None
    # Whether to prepend the prompt to the generated text
    return_full_text: bool = False
    # Stop generating tokens if a member of `stop_sequences` is generated
    stop: List[str] = []
    # Random sampling seed
    seed: Optional[int] = None
    # The value used to module the logits distribution.
    temperature: Optional[float] = None
    # The number of highest probability vocabulary tokens to keep for top-k-filtering.
    top_k: Optional[int] = None
    # If set to < 1, only the smallest set of most probable tokens with probabilities that add up to `top_p` or
    # higher are kept for generation.
    top_p: Optional[float] = None
    # truncate inputs tokens to the given size
    truncate: Optional[int] = None
    # Typical Decoding mass
    # See [Typical Decoding for Natural Language Generation](https://arxiv.org/abs/2202.00666) for more information
    typical_p: Optional[float] = None
    # Generate best_of sequences and return the one if the highest token logprobs
    best_of: Optional[int] = None
    # Watermarking with [A Watermark for Large Language Models](https://arxiv.org/abs/2301.10226)
    watermark: bool = False
    # Get generation details
    details: bool = False
    # Get decoder input token logprobs and ids
    decoder_input_details: bool = False
    # The number of highest probability vocabulary tokens to return as alternative tokens in the generation result
    return_k_alternatives: Optional[int] = None
    # Optional response format specification to constrain the generated text
    response_format: Optional[ResponseFormat] = None

    @model_validator(mode="after")
    def valid_adapter_id(self):
        adapter_id = self.adapter_id
        merged_adapters = self.merged_adapters
        if adapter_id is not None and merged_adapters is not None:
            raise ValidationError(
                "you must specify at most one of `adapter_id` or `merged_adapters`",
            )
        return self

    @field_validator("adapter_source")
    @classmethod
    def valid_adapter_source(cls, v):
        if v is not None and v not in ADAPTER_SOURCES:
            raise ValidationError(
                f"`adapter_source={v}` must be one of {ADAPTER_SOURCES}",
            )
        return v

    @field_validator("best_of")
    @classmethod
    def valid_best_of(cls, field_value, values):
        if field_value is not None:
            if field_value <= 0:
                raise ValidationError("`best_of` must be strictly positive")
            if field_value > 1 and values.data["seed"] is not None:
                raise ValidationError("`seed` must not be set when `best_of` is > 1")
            sampling = (
                values.data["do_sample"]
                | (values.data["temperature"] is not None)
                | (values.data["top_k"] is not None)
                | (values.data["top_p"] is not None)
                | (values.data["typical_p"] is not None)
            )
            if field_value > 1 and not sampling:
                raise ValidationError("you must use sampling when `best_of` is > 1")

        return field_value

    @field_validator("repetition_penalty")
    @classmethod
    def valid_repetition_penalty(cls, v):
        if v is not None and v <= 0:
            raise ValidationError("`repetition_penalty` must be strictly positive")
        return v

    @field_validator("seed")
    @classmethod
    def valid_seed(cls, v):
        if v is not None and v < 0:
            raise ValidationError("`seed` must be positive")
        return v

    @field_validator("temperature")
    @classmethod
    def valid_temp(cls, v):
        if v is not None and v < 0:
            raise ValidationError("`temperature` must be non-negative")
        return v

    @field_validator("top_k")
    @classmethod
    def valid_top_k(cls, v):
        if v is not None and v <= 0:
            raise ValidationError("`top_k` must be strictly positive")
        return v

    @field_validator("top_p")
    @classmethod
    def valid_top_p(cls, v):
        if v is not None and (v <= 0 or v >= 1.0):
            raise ValidationError("`top_p` must be > 0.0 and < 1.0")
        return v

    @field_validator("truncate")
    @classmethod
    def valid_truncate(cls, v):
        if v is not None and v <= 0:
            raise ValidationError("`truncate` must be strictly positive")
        return v

    @field_validator("typical_p")
    @classmethod
    def valid_typical_p(cls, v):
        if v is not None and (v <= 0 or v >= 1.0):
            raise ValidationError("`typical_p` must be > 0.0 and < 1.0")
        return v

    @field_validator("return_k_alternatives")
    @classmethod
    def valid_return_k_alternatives(cls, v):
        if v is not None and v <= 0:
            raise ValidationError("`return_k_alternatives` must be strictly positive")
        return v


class Request(BaseModel):
    # Prompt
    inputs: str
    # Generation parameters
    parameters: Optional[Parameters] = None
    # Whether to stream output tokens
    stream: bool = False

    @field_validator("inputs")
    @classmethod
    def valid_input(cls, v):
        if not v:
            raise ValidationError("`inputs` cannot be empty")
        return v

    @field_validator("stream")
    @classmethod
    def valid_best_of_stream(cls, field_value, values):
        parameters = values.data["parameters"]
        if parameters is not None and parameters.best_of is not None and parameters.best_of > 1 and field_value:
            raise ValidationError(
                "`best_of` != 1 is not supported when `stream` == True",
            )
        return field_value


class BatchRequest(BaseModel):
    # Prompt
    inputs: List[str]
    # Generation parameters
    parameters: Optional[Parameters] = None
    # Whether to stream output tokens
    stream: bool = False

    @field_validator("inputs")
    @classmethod
    def valid_input(cls, v):
        if not v:
            raise ValidationError("`inputs` cannot be empty")
        return v

    @field_validator("stream")
    @classmethod
    def valid_best_of_stream(cls, field_value, values):
        parameters = values.data["parameters"]
        if parameters is not None and parameters.best_of is not None and parameters.best_of > 1 and field_value:
            raise ValidationError(
                "`best_of` != 1 is not supported when `stream` == True",
            )
        return field_value


# Decoder input tokens
class InputToken(BaseModel):
    # Token ID from the model tokenizer
    id: int
    # Token text
    text: str
    # Logprob
    # Optional since the logprob of the first token cannot be computed
    logprob: Optional[float] = None


# Alternative Tokens
class AlternativeToken(BaseModel):
    # Token ID from the model tokenizer
    id: int
    # Token text
    text: str
    # Logprob
    logprob: float


# Generated tokens
class Token(BaseModel):
    # Token ID from the model tokenizer
    id: int
    # Token text
    text: str
    # Logprob
    logprob: Optional[float]
    # Is the token a special token
    # Can be used to ignore tokens when concatenating
    special: bool
    # Alternative tokens
    alternative_tokens: Optional[List[AlternativeToken]] = None
    # If token was skipped due to speculative decoding
    skipped: Optional[bool] = False


# Generation finish reason
class FinishReason(str, Enum):
    # number of generated tokens == `max_new_tokens`
    Length = "length"
    # the model generated its end of sequence token
    EndOfSequenceToken = "eos_token"
    # the model generated a text included in `stop_sequences`
    StopSequence = "stop_sequence"


# Additional sequences when using the `best_of` parameter
class BestOfSequence(BaseModel):
    # Generated text
    generated_text: str
    # Generation finish reason
    finish_reason: FinishReason
    # Number of generated tokens
    generated_tokens: int
    # Sampling seed if sampling was activated
    seed: Optional[int] = None
    # Decoder input tokens, empty if decoder_input_details is False
    prefill: List[InputToken]
    # Generated tokens
    tokens: List[Token]


# `generate` details
class Details(BaseModel):
    # Generation finish reason
    finish_reason: FinishReason
    # Number of prompt tokens
    prompt_tokens: int
    # Number of generated tokens
    generated_tokens: int
    # Number of skipped tokens
    skipped_tokens: Optional[int] = 0
    # Sampling seed if sampling was activated
    seed: Optional[int] = None
    # Decoder input tokens, empty if decoder_input_details is False
    prefill: List[InputToken]
    # Generated tokens
    tokens: List[Token]
    # Additional sequences when using the `best_of` parameter
    best_of_sequences: Optional[List[BestOfSequence]] = None


# `generate` return value
class Response(BaseModel):
    # Generated text
    generated_text: str
    # Generation details
    details: Optional[Details] = None


# `generate_stream` details
class StreamDetails(BaseModel):
    # Generation finish reason
    finish_reason: FinishReason
    # Number of prompt tokens
    prompt_tokens: int
    # Number of generated tokens
    generated_tokens: int
    # Sampling seed if sampling was activated
    seed: Optional[int] = None


# `generate_stream` return value
class StreamResponse(BaseModel):
    # Generated token
    token: Token
    # Complete generated text
    # Only available when the generation is finished
    generated_text: Optional[str] = None
    # Generation details
    # Only available when the generation is finished
    details: Optional[StreamDetails] = None


# Inference API currently deployed model
class DeployedModel(BaseModel):
    model_id: str
    sha: str
    # Suppress pydantic warning over `model_id` field
    model_config = ConfigDict(protected_namespaces=())


class EmbedResponse(BaseModel):
    # Embeddings
    embeddings: Optional[List[float]]


# Taken directly from vllm/entrypoints/openai/protocol.py
class OpenAIBaseModel(BaseModel):
    # OpenAI API does allow extra fields
    model_config = ConfigDict(extra="allow")

    # Cache class field names
    field_names: ClassVar[Optional[Set[str]]] = None

    # Predibase API token used for the request
    api_token: str = ""


class PromptTokenUsageInfo(OpenAIBaseModel):
    cached_tokens: Optional[int] = None


class UsageInfo(OpenAIBaseModel):
    prompt_tokens: int = 0
    total_tokens: int = 0
    completion_tokens: Optional[int] = 0
    prompt_token_details: Optional[List[PromptTokenUsageInfo]] = None


def random_uuid() -> str:
    return str(uuid.uuid4().hex)


class ClassificationRequest(OpenAIBaseModel):
    model: Optional[str] = None
    input: Union[List[str], str]


class ClassificationData(OpenAIBaseModel):
    index: int
    label: Optional[str]
    probs: List[float]
    num_classes: int

    def model_post_init(self, __context):
        # Remove 'api_token' field if present in the instance
        if hasattr(self, "api_token"):
            delattr(self, "api_token")


class ClassificationResponse(OpenAIBaseModel):
    id: str = Field(default_factory=lambda: f"classify-{random_uuid()}")
    object: str = "list"
    created: int = Field(default_factory=lambda: int(time.time()))
    model: str
    data: List[ClassificationData]
    # usage: UsageInfo

    def model_post_init(self, __context):
        # Remove 'usage' field if present in the instance
        if hasattr(self, "usage"):
            delattr(self, "usage")
