import time
import uuid
from enum import Enum
from typing import Any, Literal, List, Optional

from openai.types import CompletionUsage
from openai.types.chat import ChatCompletionChunk
from openai.types.chat.chat_completion import ChoiceLogprobs
from openai.types.chat.chat_completion_chunk import ChoiceDelta, Choice, ChoiceDeltaFunctionCall, \
    ChoiceDeltaToolCall
from openkosmos_ai.common.model import AIBaseConfig, LogMode
from openkosmos_core.auth import AuthTokenConfig
from pydantic import BaseModel


class Role(str, Enum):
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    FUNCTION = "function"
    TOOL = "tool"
    DEVELOPER = "developer"


class FinishReason(str, Enum):
    STOP = "stop"
    LENGTH = "length"
    TOOL_CALLS = "tool_calls"
    CONTENT_FILTER = "content_filter"
    FUNCTION_CALL = "function_call"


class FunctionCall(BaseModel):
    name: str
    arguments: str


class ChatMessage(BaseModel):
    role: Role
    content: str | None = None
    name: str | None = None
    function_call: FunctionCall | None = None


class ToolCallFunction(BaseModel):
    name: str
    arguments: str


class ToolCall(BaseModel):
    id: str
    type: Literal["function"] = "function"
    function: ToolCallFunction


class ChatCompletionRequestMessage(BaseModel):
    role: Role
    content: str | None = None
    name: str | None = None
    function_call: FunctionCall | None = None
    tool_calls: list[ToolCall] | None = None


class FunctionDefinition(BaseModel):
    name: str
    description: str | None = None
    parameters: dict[str, Any] | None = None


class ToolFunction(BaseModel):
    function: FunctionDefinition


class Tool(BaseModel):
    type: Literal["function"] = "function"
    function: FunctionDefinition


class ChatCompletionRequest(BaseModel):
    model: str
    messages: list[ChatCompletionRequestMessage]
    temperature: float | None = 0.7
    top_p: float | None = 1.0
    n: int | None = 1
    stream: bool | None = False
    stop: str | list[str] | None = None
    max_tokens: int | None = None
    presence_penalty: float | None = 0.0
    frequency_penalty: float | None = 0.0
    logit_bias: dict[str, float] | None = None
    user: str | None = None
    functions: list[FunctionDefinition] | None = None
    function_call: str | FunctionCall | None = None
    tools: list[Tool] | None = None
    tool_choice: Any | None = None

    def last_message(self):
        size = len(self.messages)
        return self.messages[size - 1] if size > 0 else None


###


class ResponseConfig(BaseModel):
    exclude_none: bool = True


class GraphProcessorConfig(AIBaseConfig):
    tags: Optional[list[str]] = ["default"]
    auth: Optional[AuthTokenConfig] = None
    response_config: ResponseConfig = ResponseConfig()
    log_mode: LogMode = LogMode.NONE


###

class ModelInfo(BaseModel):
    id: str
    object: str = "model"
    created: int = int(time.time())


class ModelInfoList(BaseModel):
    object: str = "list"
    data: list[ModelInfo] = []


###


class ChatChoiceDelta(ChoiceDelta):
    phase: Optional[str] = None
    phase_info: Optional[list] = None
    status: Literal["pending", "ongoing", "completed"] = None

    content: Optional[str] = ""
    function_call: Optional[ChoiceDeltaFunctionCall] = None
    refusal: Optional[str] = None
    role: Optional[Role] = Role.ASSISTANT
    tool_calls: Optional[List[ChoiceDeltaToolCall]] = None


class ChatChoice(Choice):
    delta: ChatChoiceDelta = ChatChoiceDelta()
    finish_reason: Optional[FinishReason] = None
    index: int = 0
    logprobs: Optional[ChoiceLogprobs] = None


class ChatCompletionResponse(ChatCompletionChunk):
    id: Optional[str] = f"chatcmpl-{str(uuid.uuid4())}"
    choices: List[ChatChoice] = [ChatChoice()]
    created: int = int(time.time())
    model: str
    object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
    service_tier: Optional[Literal["auto", "default", "flex", "scale", "priority"]] = None
    system_fingerprint: Optional[str] = None
    usage: Optional[CompletionUsage] = None

    def default_content(self, content: str):
        self.choices[0].delta.content = content
        return self

    def default_finish_reason(self, finish_reason: FinishReason):
        self.choices[0].finish_reason = finish_reason
        return self

    def default_phase(self, phase: str):
        self.choices[0].delta.phase = phase
        return self

    def get_message(self, finish_reason=None, content=None, phase=None, with_done=False,
                    response_config: ResponseConfig = ResponseConfig()):
        copy_instance = self.model_copy()
        if finish_reason is not None:
            copy_instance.default_finish_reason(FinishReason.STOP)
        if content is not None:
            copy_instance.default_content(content)
        if phase is not None:
            copy_instance.default_phase(phase)
        if with_done:
            return f"data: {copy_instance.model_dump_json(exclude_none=response_config.exclude_none)} \n\ndata: [DONE] \n\n"
        else:
            return f"data: {copy_instance.model_dump_json(exclude_none=response_config.exclude_none)} \n\n"
