# NEON AI (TM) SOFTWARE, Software Development Kit & Application Development System
# All trademark and other rights reserved by their respective owners
# Copyright 2008-2024 Neongecko.com Inc.
# BSD-3
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from this
#    software without specific prior written permission.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
# CONTRIBUTORS  BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA,
# OR PROFITS;  OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE,  EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

from typing import List, Optional, Any, Dict, Literal, Union
from pydantic import Field

from neon_data_models.models.base.contexts import MQContext
from neon_data_models.models.api.llm import BrainForgeLLM, LLMRequest, LLMResponse, LLMPersona
from neon_data_models.models.api.http.brainforge import LLMGetModelsHttpResponse, LLMGetPersonasHttpRequest


class LLMGetModels(MQContext):
    user_id: str = Field(
        description="ID of user to get models for")


class LLMGetModelsResponse(MQContext, LLMGetModelsHttpResponse):
    pass


class LLMGetPersonas(LLMGetModels, LLMGetPersonasHttpRequest):
    @property
    def model_name(self):
        return self.model_id.split("@")[0]

    @property
    def model_version(self):
        return self.model_id.split("@", 1)[1]


class LLMGetPersonasResponse(MQContext):
    model: Optional[BrainForgeLLM] = Field(
        "Full configuration of requested model if model is loaded and access "
        "is allowed, else None.")

    @property
    def personas(self) -> List[LLMPersona]:
        """
        Convenience property defined to easily reference the personas requested
        """
        return self.model.personas if self.model else []


class LLMGetInference(LLMRequest, MQContext):
    user_id: str = Field("ID of user making the request")

    def as_llm_request(self):
        """
        Get a plain `LLMRequest` object from this `LLMGetInference` object.
        """
        return LLMRequest(**self.model_dump())


class LLMGetCompletion(LLMGetModels):
    model: str = Field(description="Model to request (<name>@<revision>)")
    completion_kwargs: Dict[str, Any] = Field(
        description="Dictionary of kwargs to pass to OpenAI completion request")


class LLMGetCompletionResponse(MQContext):
    openai_response: dict = Field(
        description="OpenAI ChatCompletion model")


class LLMGetTokenizerChatTemplatedString(LLMGetModels):
    model: str = Field(description="Model to request (<name>@<revision>)")
    messages: List[Dict[Literal["role", "content"], str]] = Field(
        description="List of dict messages in OpenAI format")
    add_generation_prompt: bool = Field(
        description="If true, assistant start tokens will be appended to the "
                    "formatted output.")
    tokenize: bool = Field(
        False,
        description="If true, a list of token strings is returned, "
                    "else a single string")


class LLMGetTokenizerChatTemplatedStringResponse(MQContext):
    prompt: Union[List[str], str] = Field(
        description="Prompt generated by the tokenizer")


class LLMGetInferenceResponse(LLMResponse, MQContext):
    pass


__all__ = [LLMGetModels.__name__, LLMGetModelsResponse.__name__,
           LLMGetPersonas.__name__, LLMGetPersonasResponse.__name__,
           LLMGetInference.__name__, LLMGetInferenceResponse.__name__,
           LLMGetCompletion.__name__, LLMGetCompletionResponse.__name__,
           LLMGetTokenizerChatTemplatedString.__name__,
           LLMGetTokenizerChatTemplatedStringResponse.__name__]
