from __future__ import annotations

import json
import time
import typing
from collections.abc import Iterable
from copy import deepcopy
from pathlib import Path
from typing import Any, Generic, Literal

from pydantic import BaseModel as B_Model
from pydantic import Field

# Pydantic models

T = typing.TypeVar("T", None, str, None | typing.Literal[""])
T_INT = typing.TypeVar("T_INT", int, None)


class BaseModel(B_Model):
    """BaseModel+dict duck typing"""

    def __str__(self) -> str:
        return json.dumps(self.model_dump(), ensure_ascii=True)

    def __repr__(self) -> str:
        return self.__str__()

    def __getitem__(self, key: str) -> Any:
        return self.model_dump()[key]

    def __setitem__(self, key: str, value: Any) -> None:
        self.__setattr__(key, value)


class ModelConfig(BaseModel):
    """Model configuration"""

    top_k: int = Field(
        default=50,
        description="TopK (Some model adapters may not support this parameter)",
    )
    top_p: float = Field(default=0.8, description="TopP")
    temperature: float = Field(default=0.6, description="Temperature")
    stream: bool = Field(
        default=False,
        description="Whether to enable streaming response (output by character)",
    )
    multimodal: bool = Field(
        default=False,
        description="Whether to support multimodal input (e.g. image recognition)",
    )


class ModelPreset(BaseModel):
    model: str = Field(
        default="", description="Name of the AI model to use (e.g. gpt-3.5-turbo)"
    )
    name: str = Field(
        default="default", description="Identifier name for current preset"
    )
    base_url: str = Field(
        default="",
        description="Base address of API service (use OpenAI default if empty)",
    )
    api_key: str = Field(default="", description="Key required to access API")
    protocol: str = Field(default="__main__", description="Protocol adapter type")
    config: ModelConfig = Field(default_factory=ModelConfig)
    extra: dict[str, Any] = Field(default_factory=dict)

    @classmethod
    def load(cls, path: Path):
        if path.exists():
            with path.open(
                "r",
                encoding="utf-8",
            ) as f:
                data = json.load(f)
            return cls.model_validate(data)
        return cls()  # Return default values

    def save(self, path: Path):
        with path.open("w", encoding="u8") as f:
            json.dump(self.model_dump(), f, indent=4, ensure_ascii=False)


class Function(BaseModel):
    arguments: str
    """
    The arguments to call the function with, as generated by the model in JSON
    format. Note that the model does not always generate valid JSON, and may
    hallucinate parameters not defined by your function schema. Validate the
    arguments in your code before calling your function.
    """

    name: str
    """The name of the function to call."""


class ToolCall(BaseModel):
    id: str
    """The ID of the tool call."""

    function: Function
    """The function that the model called."""

    type: typing.Literal["function"] = "function"
    """The type of the tool. Currently, only `function` is supported."""


T_TOOL = typing.TypeVar("T_TOOL", list[ToolCall], None, list[ToolCall] | None)


class UniResponseUsage(BaseModel, Generic[T_INT]):
    prompt_tokens: T_INT
    completion_tokens: T_INT
    total_tokens: T_INT


class UniResponse(
    BaseModel,
    Generic[T, T_TOOL],
):
    """Unified response format"""

    role: Literal["assistant"] = Field(
        default="assistant",  # Regardless of whether there's content/tool_call, role is assistant
        description="Role",
    )

    usage: UniResponseUsage | None = None
    content: T = Field(
        ...,
        description="Content",
        exclude_if=lambda x: x is None,
    )
    tool_calls: T_TOOL = Field(
        ...,
        description="Tool call results",
        exclude_if=lambda x: x is None,
    )


class ImageUrl(BaseModel):
    url: str = Field(..., description="Image URL")


class Content(BaseModel): ...


class ImageContent(Content):
    type: Literal["image_url"] = "image_url"
    image_url: ImageUrl = Field(..., description="Image URL")


class TextContent(Content):
    type: Literal["text"] = "text"
    text: str = Field(..., description="Text content")


CT_MAP: dict[str, type[Content]] = {
    "image_url": ImageContent,
    "text": TextContent,
}

USER_INPUT = list[TextContent | ImageContent] | str

_T = typing.TypeVar(
    "_T",
    str,
    None,
    list[TextContent],
    list[TextContent | ImageContent],
    USER_INPUT,
)


class Message(BaseModel, Generic[_T]):
    role: Literal["user", "assistant", "system"] = Field(..., description="Role")
    content: _T = Field(..., description="Content")
    tool_calls: list[ToolCall] | None = Field(
        default=None, description="Tool calls", exclude_if=lambda x: x is None
    )


class ToolResult(BaseModel):
    role: Literal["tool"] = Field(..., description="Role")
    name: str = Field(..., description="Tool name")
    content: str = Field(..., description="Tool return content")
    tool_call_id: str = Field(..., description="Tool call ID")


class MemoryModel(BaseModel):
    messages: list[CONTENT_LIST_TYPE_ITEM] = Field(default_factory=list)
    time: float = Field(default_factory=time.time, description="Timestamp")
    abstract: str = Field(default="", description="Summary")


CONTENT_LIST_TYPE_ITEM = Message | ToolResult
CONTENT_LIST_TYPE = list[CONTENT_LIST_TYPE_ITEM]


class SendMessageWrap(Iterable[CONTENT_LIST_TYPE_ITEM]):
    """Wrapper class for CONTENT_LIST_TYPE"""

    train: Message[str]  # system message
    memory: CONTENT_LIST_TYPE  # Messages without system message
    user_query: Message
    end_messages: CONTENT_LIST_TYPE  # End messages

    def __init__(
        self,
        train: dict[str, str] | Message[str],
        memory: CONTENT_LIST_TYPE | MemoryModel,
        user_query: Message | None = None,
    ):
        self.train = (
            train if isinstance(train, Message) else Message.model_validate(train)
        )
        self.end_messages = []
        self.memory = memory if isinstance(memory, list) else memory.messages
        query = user_query or self.memory[-1]
        if isinstance(query, ToolResult) or query.role != "user":
            raise ValueError("Invalid query message, expecting user message!")
        self.user_query = query
        if not user_query:
            self.memory.pop()

    @classmethod
    def validate_messages(cls, messages: CONTENT_LIST_TYPE) -> SendMessageWrap:
        train = messages[0]
        if train.role != "system":  # Fall back to match the first system message
            for idx, msg in enumerate(messages):
                if msg.role == "system":
                    train = msg
                    messages.pop(idx)
                    memory = messages
                    break
            else:
                raise ValueError("Invalid messages, expecting system message!")
        else:
            memory = messages[1:]
        return cls(
            train,
            memory,
        )

    def __len__(self) -> int:
        return len(self.memory) + 2 + len(self.end_messages)

    def __iter__(self) -> typing.Iterator[CONTENT_LIST_TYPE_ITEM]:
        yield self.train
        yield from self.memory
        yield self.user_query
        yield from self.end_messages

    def copy(self) -> SendMessageWrap:
        return deepcopy(self)

    def unwrap(self, exclude_system: bool = False) -> CONTENT_LIST_TYPE:
        system_msg: CONTENT_LIST_TYPE = [self.train] if not exclude_system else []
        return [*system_msg, *self.memory, self.user_query, *self.end_messages]

    def get_train(self) -> Message[str]:
        return self.train

    def get_memory(self) -> CONTENT_LIST_TYPE:
        return self.memory

    def get_user_query(self) -> Message:
        return self.user_query

    def append(self, message: CONTENT_LIST_TYPE_ITEM) -> None:
        self.end_messages.append(message)

    def extend(self, messages: CONTENT_LIST_TYPE) -> None:
        self.end_messages.extend(messages)
