from __future__ import annotations

import contextlib
import json
import time
import typing
from collections.abc import Iterable, Sequence
from copy import deepcopy
from datetime import datetime, timedelta
from typing import Any, Generic, Literal, overload

from nonebot_plugin_orm import AsyncSession, Model, get_session
from pydantic import BaseModel as B_Model
from pydantic import Field
from sqlalchemy import (
    JSON,
    BigInteger,
    Boolean,
    DateTime,
    Float,
    ForeignKey,
    Index,
    Integer,
    String,
    Text,
    UniqueConstraint,
    delete,
    insert,
    select,
    text,
    update,
)
from sqlalchemy.orm import Mapped, mapped_column
from typing_extensions import Self

from ..config import config_manager
from .lock import database_lock

# Pydantic 模型

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


class NoActionContext:
    """NoActionContext，仅作占位符，用于实现 with 语法"""

    async def __aenter__(self) -> Self:
        """__aenter__"""
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """__aexit__"""
        pass

    def __enter__(self) -> Self:
        """__enter__"""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        """__exit__"""
        pass


class BaseModel(B_Model):
    """BaseModel+dict鸭子类型"""

    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 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],
):
    """统一响应格式"""

    role: Literal["assistant", "function"] = "assistant"
    usage: UniResponseUsage | None = None
    content: T = Field(
        ...,
        description="内容",
        exclude_if=lambda x: x is None,
    )
    tool_calls: T_TOOL = Field(
        ...,
        description="工具调用结果",
        exclude_if=lambda x: x is None,
    )


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


class Content(BaseModel): ...


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


class TextContent(Content):
    type: Literal["text"] = "text"
    text: str = Field(..., description="文本内容")


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

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


class Message(BaseModel, Generic[_T]):
    role: Literal["user", "assistant", "system"] = Field(
        default="assistant", description="角色"
    )
    content: _T = Field(..., description="内容")
    tool_calls: list[ToolCall] | None = Field(
        default=None, description="工具调用", exclude_if=lambda x: x is None
    )


class ToolResult(BaseModel):
    role: Literal["tool"] = Field(default="tool", description="角色")
    name: str = Field(..., description="工具名称")
    content: str = Field(..., description="工具返回内容")
    tool_call_id: str = Field(..., description="工具调用ID")


class MemoryModel(BaseModel):
    messages: list[SEND_MESSAGES_ITEM] = Field(default_factory=list)
    time: float = Field(default_factory=time.time, description="时间戳")
    abstract: str = Field(default="", description="摘要")


class SessionMemoryModel(MemoryModel):
    id: int | None = Field(None, description="会话在数据库的ID")

    def __init__(self, **data):
        super().__init__(**data)
        self.__dict__["_dirty"] = False
        if "messages" in data:
            self.__dict__["_dirty"] = True

    def __setattr__(self, name, value):
        if name == "messages":
            self.__dict__["_dirty"] = True
        object.__setattr__(self, name, value)

    def __getattribute__(self, name):
        # 如果获取的是messages属性，则标记为dirty
        if name == "messages":
            value = object.__getattribute__(self, name)
            with contextlib.suppress(Exception):
                object.__setattr__(self, "_dirty", True)

            return value
        else:
            return object.__getattribute__(self, name)

    @property
    def __dirty__(self):
        """获取dirty状态"""
        return self.__dict__.get("_dirty", False)

    @__dirty__.setter
    def __dirty__(self, value: bool):
        """设置dirty状态"""
        object.__setattr__(self, "_dirty", value)

    async def delete(self, arg_session: AsyncSession | None = None):
        if self.id is None:
            raise ValueError("无法删除未保存的模型")

        async with database_lock(self.id):
            session = arg_session or get_session()

            stmt = delete(MemorySessions).where(MemorySessions.id == self.id)
            async with NoActionContext() if arg_session else session:
                await session.execute(stmt)
            if not arg_session:
                await session.commit()

    async def save(
        self,
        ins_id: int = 0,
        is_group: bool = False,
        arg_session: AsyncSession | None = None,
    ) -> None:
        # 只有在dirty状态下才保存
        if not self.__dirty__:
            return
        session = arg_session or get_session()
        async with NoActionContext() if arg_session else session:
            if self.id is not None:
                stmt = (
                    update(MemorySessions)
                    .where(MemorySessions.id == self.id)
                    .values(data=self.model_dump(exclude={"id"}))
                )
                await session.execute(stmt)
            elif ins_id:
                session.add(
                    MemorySessions(
                        data=self.model_dump(exclude={"id"}),
                        created_at=time.time(),
                        is_group=is_group,
                        ins_id=ins_id,
                    )
                )
            else:
                raise ValueError(
                    "Neither id nor ins_id provided; cannot persist memory session"
                )
            if not arg_session:
                await session.commit()

        # 保存后重置dirty标志
        self.__dirty__ = False


SEND_MESSAGES_ITEM = Message | ToolResult
SEND_MESSAGES = list[SEND_MESSAGES_ITEM]


class SendMessageWrap(Iterable[SEND_MESSAGES_ITEM]):
    """SEND_MESSAGES的包装类"""

    train: Message  # system 消息
    memory: SEND_MESSAGES  # 无system消息的消息

    def __init__(
        self, train: dict[str, str] | Message[str], memory: SEND_MESSAGES | MemoryModel
    ):
        self.train = (
            train if isinstance(train, Message) else Message.model_validate(train)
        )
        self.memory = memory if isinstance(memory, list) else memory.messages

    @classmethod
    def validate_messages(cls, messages: SEND_MESSAGES) -> 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) + 1

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

    def copy(self) -> SendMessageWrap:
        return SendMessageWrap(deepcopy(self.train), deepcopy(self.memory))

    def unwrap(self) -> SEND_MESSAGES:
        return [self.train, *self.memory]

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

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


class InsightsModel(BaseModel):
    date: str = Field(
        default_factory=lambda: datetime.now().strftime("%Y-%m-%d"), description="日期"
    )
    token_input: int = Field(..., description="输入token使用量")
    token_output: int = Field(..., description="输出token使用量")
    usage_count: int = Field(..., description="聊天请求次数")

    @classmethod
    async def get_all(cls) -> list[Self]:
        async with database_lock():
            async with get_session() as session:
                await cls._delete_expired(
                    days=config_manager.config.usage_limit.global_insights_expire_days,
                    session=session,
                )
                stmt = select(GlobalInsights)
                insights = (await session.execute(stmt)).scalars().all()
                session.add_all(insights)
                return [cls.model_validate(x, from_attributes=True) for x in insights]

    @classmethod
    async def get(cls) -> Self:
        date_now = datetime.now().strftime("%Y-%m-%d")
        async with database_lock(date_now):
            async with get_session() as session:
                await cls._delete_expired(
                    days=config_manager.config.usage_limit.global_insights_expire_days,
                    session=session,
                )
                if (
                    insights := (
                        await session.execute(
                            select(GlobalInsights).where(
                                GlobalInsights.date == date_now
                            )
                        )
                    ).scalar_one_or_none()
                ) is None:
                    stmt = insert(GlobalInsights).values(date=date_now)
                    await session.execute(stmt)
                    insights = (
                        await session.execute(
                            select(GlobalInsights).where(
                                GlobalInsights.date == date_now
                            )
                        )
                    ).scalar_one()
                session.add(insights)
                instance = cls.model_validate(insights, from_attributes=True)
            return instance

    async def save(self):
        """保存数据"""
        async with database_lock(self.date):
            async with get_session() as session:
                await self._delete_expired(
                    days=config_manager.config.usage_limit.global_insights_expire_days,
                    session=session,
                )
                stmt = select(GlobalInsights).where(GlobalInsights.date == self.date)
                if ((await session.execute(stmt)).scalar_one_or_none()) is None:
                    stmt = insert(GlobalInsights).values(
                        **{
                            k: v
                            for k, v in self.model_dump().items()
                            if hasattr(GlobalInsights, k)
                        }
                    )
                    await session.execute(stmt)
                    await session.commit()
                else:
                    stmt = (
                        update(GlobalInsights)
                        .where(GlobalInsights.date == self.date)
                        .values(
                            **{
                                k: v
                                for k, v in self.model_dump().items()
                                if hasattr(GlobalInsights, k)
                            }
                        )
                    )
                    await session.execute(stmt)
                    await session.commit()

    @staticmethod
    async def _delete_expired(*, days: int, session: AsyncSession) -> None:
        """
        删除过期的记录

        Args:
            days: 保留天数，超过此天数的记录将被删除
        """
        # 计算截止日期
        cutoff_date = datetime.now() - timedelta(days=days)

        # 删除过期记录
        stmt = delete(GlobalInsights).where(
            GlobalInsights.date < cutoff_date.strftime("%Y-%m-%d")
        )
        await session.execute(stmt)
        await session.commit()


# Sqlalchemy 模型


class GlobalInsights(Model):
    __tablename__ = "suggarchat_global_insights"
    date: Mapped[str] = mapped_column(
        String(64),
        primary_key=True,
        default=lambda: datetime.now().strftime("%Y-%m-%d"),
    )
    token_input: Mapped[int] = mapped_column(
        BigInteger, default=0, server_default=text("0")
    )
    token_output: Mapped[int] = mapped_column(
        BigInteger, default=0, server_default=text("0")
    )
    usage_count: Mapped[int] = mapped_column(Integer, default=0)


class MemorySessions(Model):
    __tablename__ = "suggarchat_memory_sessions"
    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    ins_id: Mapped[int] = mapped_column(
        ForeignKey("suggarchat_memory_data.ins_id"), nullable=False
    )
    is_group: Mapped[bool] = mapped_column(
        ForeignKey("suggarchat_memory_data.is_group"),
        nullable=False,
        default=False,
    )
    created_at: Mapped[float] = mapped_column(Float, nullable=False)
    data: Mapped[dict[str, Any]] = mapped_column(
        JSON, nullable=False, server_default=text("'{}'")
    )
    __table_args__ = (
        Index("idx_sessions_ins_id", "ins_id"),
        Index("idx_sessions_is_group", "is_group"),
        Index("idx_sessions_created_at", "created_at"),
    )

    @classmethod
    async def _expire(cls, ins_id: int, is_group: bool, keep_count: int = 20):
        """
        保留特定数量的sessions，移除多余的会话记录

        Args:
            session: 数据库会话
            ins_id: 实例ID
            is_group: 是否为群组
            keep_count: 要保留的会话数量，默认为20
        """
        # 查询指定ins_id和is_group的所有会话，按创建时间倒序排列
        async with get_session() as session:
            stmt = (
                select(cls.id)
                .where(cls.ins_id == ins_id, cls.is_group == is_group)
                .order_by(cls.created_at.desc())
                .offset(keep_count)  # 跳过要保留的数量，获取需要删除的记录
            )

            result = await session.execute(stmt)
            ids_to_delete = [row[0] for row in result.fetchall()]

            if ids_to_delete:
                # 删除超过保留数量的会话记录
                delete_stmt = delete(cls).where(cls.id.in_(ids_to_delete))
                await session.execute(delete_stmt)

                # 提交更改
                await session.commit()

    @classmethod
    async def get(
        cls, session: AsyncSession, ins_id: int, is_group: bool
    ) -> Sequence[Self]:
        async with database_lock(ins_id, is_group):
            await cls._expire(ins_id, is_group)
            stmt = select(cls).where(cls.ins_id == ins_id, cls.is_group == is_group)
            data = (await session.execute(stmt)).scalars().all()
            session.add_all(data)
            return data


class Memory(Model):
    __tablename__ = "suggarchat_memory_data"
    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
    ins_id: Mapped[int] = mapped_column(BigInteger, nullable=False)
    is_group: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
    memory_json: Mapped[dict[str, Any]] = mapped_column(
        JSON,
        default=MemoryModel().model_dump(),
        nullable=False,
        server_default=text("'{}'"),
    )
    time: Mapped[datetime] = mapped_column(
        DateTime, default=datetime.now, nullable=False
    )
    usage_count: Mapped[int] = mapped_column(Integer, default=0)
    input_token_usage: Mapped[int] = mapped_column(BigInteger, default=0)
    output_token_usage: Mapped[int] = mapped_column(BigInteger, default=0)
    __table_args__ = (
        UniqueConstraint("ins_id", "is_group", name="uq_ins_id_is_group"),
        Index("idx_ins_id", "ins_id"),
        Index("idx_is_group", "is_group"),
    )


class GroupConfig(Model):
    __tablename__ = "suggarchat_group_config"
    id: Mapped[int] = mapped_column(
        Integer,
        primary_key=True,
        autoincrement=True,
    )
    group_id: Mapped[int] = mapped_column(
        BigInteger,
        ForeignKey("suggarchat_memory_data.ins_id"),
        nullable=False,
    )
    enable: Mapped[bool] = mapped_column(Boolean, default=True)
    prompt: Mapped[str] = mapped_column(Text, default="")
    fake_people: Mapped[bool] = mapped_column(Boolean, default=False)
    last_updated: Mapped[datetime] = mapped_column(DateTime, default=datetime.now)

    __table_args__ = (
        UniqueConstraint("group_id", name="uq_suggarchat_config_group_id"),
        Index("idx_suggarchat_group_id", "group_id"),
    )


@overload
async def get_or_create_data(
    *, session: AsyncSession, ins_id: int, for_update: bool = False
) -> Memory: ...
@overload
async def get_or_create_data(
    *,
    session: AsyncSession,
    ins_id: int,
    is_group: bool = True,
    for_update: bool = False,
) -> tuple[GroupConfig, Memory]: ...


async def get_or_create_data(
    *,
    session: AsyncSession,
    ins_id: int,
    is_group: bool = False,
    for_update: bool = False,
) -> Memory | tuple[GroupConfig, Memory]:
    async with database_lock(ins_id, is_group):
        stmt = select(Memory).where(
            Memory.ins_id == ins_id, Memory.is_group == is_group
        )
        stmt = stmt.with_for_update() if for_update else stmt
        result = await session.execute(stmt)
        if not (memory := result.scalar_one_or_none()):
            stmt = insert(Memory).values(ins_id=ins_id, is_group=is_group)
            await session.execute(stmt)
            await session.commit()
            stmt = select(Memory).where(
                Memory.ins_id == ins_id, Memory.is_group == is_group
            )
            stmt = stmt.with_for_update() if for_update else stmt
            memory = (await session.execute(stmt)).scalar_one()
        session.add(memory)
        if not is_group:
            return memory
        stmt = select(GroupConfig).where(GroupConfig.group_id == ins_id)
        stmt = stmt.with_for_update() if for_update else stmt
        result = await session.execute(stmt)
        if not (group_config := result.scalar_one_or_none()):
            stmt = insert(GroupConfig).values(group_id=ins_id)
            await session.execute(stmt)
            await session.commit()
            stmt = select(GroupConfig).where(GroupConfig.group_id == ins_id)
            group_config = (await session.execute(stmt)).scalar_one()
        session.add(group_config)
        return group_config, memory
