"""
对话域持久化 Repository（Async）
"""

import uuid as uuid_lib

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload
from sqlalchemy.orm.attributes import flag_modified

from yuxi.storage.postgres.models_business import Conversation, ConversationStats, Message, ToolCall
from yuxi.utils import logger
from yuxi.utils.datetime_utils import utc_now_naive

MAX_CONVERSATION_TITLE_LENGTH = 255


class ConversationRepository:
    def __init__(self, db_session: AsyncSession):
        self.db = db_session

    def _normalize_title(self, title: str | None) -> str | None:
        if title is None:
            return None
        normalized = str(title).strip()
        if not normalized:
            return ""
        if len(normalized) > MAX_CONVERSATION_TITLE_LENGTH:
            logger.warning(
                f"Conversation title too long ({len(normalized)}), truncate to {MAX_CONVERSATION_TITLE_LENGTH}"
            )
            return normalized[:MAX_CONVERSATION_TITLE_LENGTH]
        return normalized

    async def create_conversation(
        self,
        uid: str,
        agent_id: str,
        title: str | None = None,
        thread_id: str | None = None,
        metadata: dict | None = None,
    ) -> Conversation:
        if not thread_id:
            thread_id = str(uuid_lib.uuid4())

        metadata = (metadata or {}).copy()
        metadata.setdefault("attachments", [])

        normalized_title = self._normalize_title(title)

        conversation = Conversation(
            thread_id=thread_id,
            uid=str(uid),
            agent_id=agent_id,
            title=normalized_title or "New Conversation",
            status="active",
            extra_metadata=metadata,
        )

        self.db.add(conversation)
        await self.db.flush()

        stats = ConversationStats(conversation_id=conversation.id)
        self.db.add(stats)
        await self.db.commit()
        await self.db.refresh(conversation)

        logger.info(f"Created conversation: {conversation.thread_id} for user {uid}")
        return conversation

    async def get_conversation_by_thread_id(self, thread_id: str) -> Conversation | None:
        result = await self.db.execute(select(Conversation).where(Conversation.thread_id == thread_id))
        return result.scalar_one_or_none()

    async def _get_conversation_by_id(self, conversation_id: int) -> Conversation | None:
        result = await self.db.execute(select(Conversation).where(Conversation.id == conversation_id))
        return result.scalar_one_or_none()

    def _ensure_metadata(self, conversation: Conversation) -> dict:
        metadata = dict(conversation.extra_metadata or {})
        attachments = metadata.get("attachments", [])
        metadata["attachments"] = [dict(item) for item in attachments if isinstance(item, dict)]
        return metadata

    async def _save_metadata(self, conversation: Conversation, metadata: dict) -> None:
        conversation.extra_metadata = metadata
        flag_modified(conversation, "extra_metadata")
        conversation.updated_at = utc_now_naive()
        await self.db.commit()
        await self.db.refresh(conversation)

    async def add_message(
        self,
        conversation_id: int,
        role: str,
        content: str,
        message_type: str = "text",
        extra_metadata: dict | None = None,
        image_content: str | None = None,
        run_id: str | None = None,
        request_id: str | None = None,
        delivery_status: str = "complete",
    ) -> Message:
        message = Message(
            conversation_id=conversation_id,
            role=role,
            content=content,
            message_type=message_type,
            extra_metadata=extra_metadata or {},
            image_content=image_content,
            run_id=run_id,
            request_id=request_id,
            delivery_status=delivery_status,
        )

        self.db.add(message)
        conversation = await self._get_conversation_by_id(conversation_id)
        if conversation:
            conversation.updated_at = utc_now_naive()

        await self.db.commit()
        await self.db.refresh(message)

        await self._update_message_count(conversation_id)

        logger.debug(f"Added {role} message to conversation {conversation_id}")
        return message

    async def add_message_by_thread_id(
        self,
        thread_id: str,
        role: str,
        content: str,
        message_type: str = "text",
        extra_metadata: dict | None = None,
        image_content: str | None = None,
        run_id: str | None = None,
        request_id: str | None = None,
        delivery_status: str = "complete",
    ) -> Message | None:
        conversation = await self.get_conversation_by_thread_id(thread_id)
        if not conversation:
            logger.warning(f"Conversation not found for thread_id: {thread_id}")
            return None

        return await self.add_message(
            conversation_id=conversation.id,
            role=role,
            content=content,
            message_type=message_type,
            extra_metadata=extra_metadata,
            image_content=image_content,
            run_id=run_id,
            request_id=request_id,
            delivery_status=delivery_status,
        )

    async def add_tool_call(
        self,
        message_id: int,
        tool_name: str,
        tool_input: dict | None = None,
        tool_output: str | None = None,
        status: str = "pending",
        error_message: str | None = None,
        langgraph_tool_call_id: str | None = None,
    ) -> ToolCall:
        if langgraph_tool_call_id:
            existing = await self.get_tool_call_by_langgraph_id(langgraph_tool_call_id)
            if existing:
                logger.debug(
                    "Tool call already exists for langgraph_tool_call_id=%s, skip insert",
                    langgraph_tool_call_id,
                )
                return existing

        tool_call = ToolCall(
            message_id=message_id,
            tool_name=tool_name,
            tool_input=tool_input or {},
            tool_output=tool_output,
            status=status,
            error_message=error_message,
            langgraph_tool_call_id=langgraph_tool_call_id,
        )

        self.db.add(tool_call)
        await self.db.commit()
        await self.db.refresh(tool_call)

        logger.debug(f"Added tool call {tool_name} to message {message_id}")
        return tool_call

    async def get_messages(self, conversation_id: int, limit: int | None = None, offset: int = 0) -> list[Message]:
        query = (
            select(Message)
            .options(
                selectinload(Message.tool_calls),
                selectinload(Message.feedbacks),
            )
            .where(Message.conversation_id == conversation_id)
            .order_by(Message.created_at.asc())
        )

        if limit:
            query = query.limit(limit).offset(offset)

        result = await self.db.execute(query)
        return list(result.scalars().unique().all())

    async def get_messages_by_thread_id(
        self, thread_id: str, limit: int | None = None, offset: int = 0
    ) -> list[Message]:
        conversation = await self.get_conversation_by_thread_id(thread_id)
        if not conversation:
            logger.warning(f"Conversation not found for thread_id: {thread_id}")
            return []

        return await self.get_messages(conversation.id, limit, offset)

    async def list_conversations(
        self,
        uid: str | None = None,
        agent_id: str | None = None,
        status: str = "active",
        limit: int | None = None,
        offset: int = 0,
    ) -> list[Conversation]:
        """List conversations with pinned conversations always included first.

        The limit applies only to non-pinned conversations to ensure pinned
        conversations are always visible in the list.
        """

        base_conditions = [Conversation.status == status]
        if uid:
            base_conditions.append(Conversation.uid == str(uid))
        if agent_id:
            base_conditions.append(Conversation.agent_id == agent_id)

        # First, get all pinned conversations (no limit)
        pinned_query = (
            select(Conversation)
            .where(*base_conditions)
            .where(Conversation.is_pinned)
            .order_by(Conversation.updated_at.desc())
        )
        result = await self.db.execute(pinned_query)
        pinned_conversations = list(result.scalars().all())

        # Then, get non-pinned conversations with limit/offset
        remaining_limit = None
        remaining_offset = offset

        if limit is not None:
            # Calculate how many slots are taken by pinned conversations
            pinned_count = len(pinned_conversations)
            if pinned_count >= limit:
                # All slots taken by pinned conversations
                return pinned_conversations[:limit]
            remaining_limit = limit - pinned_count

        if remaining_limit is not None and remaining_limit > 0:
            non_pinned_query = (
                select(Conversation)
                .where(*base_conditions)
                .where(~Conversation.is_pinned)
                .order_by(Conversation.updated_at.desc())
                .limit(remaining_limit)
                .offset(remaining_offset)
            )
            result = await self.db.execute(non_pinned_query)
            non_pinned_conversations = list(result.scalars().all())
        else:
            non_pinned_conversations = []

        return pinned_conversations + non_pinned_conversations

    async def update_conversation(
        self,
        thread_id: str,
        title: str | None = None,
        status: str | None = None,
        metadata: dict | None = None,
        is_pinned: bool | None = None,
    ) -> Conversation | None:
        conversation = await self.get_conversation_by_thread_id(thread_id)
        if not conversation:
            return None

        normalized_title = self._normalize_title(title)
        if normalized_title is not None:
            conversation.title = normalized_title
        if status is not None:
            conversation.status = status
        if is_pinned is not None:
            conversation.is_pinned = is_pinned

        if metadata is not None:
            current_metadata = conversation.extra_metadata or {}
            current_metadata.update(metadata)
            conversation.extra_metadata = current_metadata

        conversation.updated_at = utc_now_naive()
        await self.db.commit()
        await self.db.refresh(conversation)

        logger.info(f"Updated conversation {thread_id}")
        return conversation

    async def delete_conversation(self, thread_id: str, soft_delete: bool = True) -> bool:
        conversation = await self.get_conversation_by_thread_id(thread_id)
        if not conversation:
            return False

        if soft_delete:
            conversation.status = "deleted"
            await self.db.commit()
            logger.info(f"Soft deleted conversation {thread_id}")
        else:
            self.db.delete(conversation)
            await self.db.commit()
            logger.info(f"Permanently deleted conversation {thread_id}")

        return True

    async def get_stats(self, conversation_id: int) -> ConversationStats | None:
        result = await self.db.execute(
            select(ConversationStats).where(ConversationStats.conversation_id == conversation_id)
        )
        return result.scalar_one_or_none()

    async def update_stats(
        self,
        conversation_id: int,
        tokens_used: int | None = None,
        model_used: str | None = None,
        user_feedback: dict | None = None,
    ) -> ConversationStats | None:
        stats = await self.get_stats(conversation_id)
        if not stats:
            return None

        if tokens_used is not None:
            stats.total_tokens += tokens_used
        if model_used is not None:
            stats.model_used = model_used
        if user_feedback is not None:
            stats.user_feedback = user_feedback

        stats.updated_at = utc_now_naive()
        await self.db.commit()
        await self.db.refresh(stats)

        return stats

    async def get_tool_call_by_langgraph_id(self, langgraph_tool_call_id: str) -> ToolCall | None:
        result = await self.db.execute(
            select(ToolCall)
            .where(ToolCall.langgraph_tool_call_id == langgraph_tool_call_id)
            .order_by(ToolCall.created_at.desc())
            .limit(1)
        )
        return result.scalar_one_or_none()

    async def update_tool_call_output(
        self,
        langgraph_tool_call_id: str,
        tool_output: str,
        status: str = "success",
        error_message: str | None = None,
    ) -> ToolCall | None:
        tool_call = await self.get_tool_call_by_langgraph_id(langgraph_tool_call_id)
        if not tool_call:
            logger.warning(f"Tool call not found for langgraph_tool_call_id: {langgraph_tool_call_id}")
            return None

        tool_call.tool_output = tool_output
        tool_call.status = status
        if error_message:
            tool_call.error_message = error_message

        await self.db.commit()
        await self.db.refresh(tool_call)

        logger.debug(f"Updated tool call {langgraph_tool_call_id} with output")
        return tool_call

    async def _update_message_count(self, conversation_id: int) -> None:
        from sqlalchemy import func

        stats = await self.get_stats(conversation_id)
        if stats:
            result = await self.db.execute(select(func.count()).where(Message.conversation_id == conversation_id))
            message_count = result.scalar()
            stats.message_count = message_count
            await self.db.commit()

    async def get_attachments(self, conversation_id: int) -> list[dict]:
        conversation = await self._get_conversation_by_id(conversation_id)
        if not conversation:
            return []
        metadata = self._ensure_metadata(conversation)
        return list(metadata.get("attachments", []))

    async def get_attachments_by_thread_id(self, thread_id: str) -> list[dict]:
        conversation = await self.get_conversation_by_thread_id(thread_id)
        if not conversation:
            return []
        return await self.get_attachments(conversation.id)

    async def add_attachment(self, conversation_id: int, attachment_info: dict) -> dict | None:
        conversation = await self._get_conversation_by_id(conversation_id)
        if not conversation:
            return None

        metadata = self._ensure_metadata(conversation)
        attachments = metadata.get("attachments", [])
        attachments = [item for item in attachments if item.get("file_id") != attachment_info.get("file_id")]
        attachments.append(attachment_info)
        metadata["attachments"] = attachments
        await self._save_metadata(conversation, metadata)
        return attachment_info

    async def add_attachments(self, conversation_id: int, attachment_infos: list[dict]) -> list[dict] | None:
        conversation = await self._get_conversation_by_id(conversation_id)
        if not conversation:
            return None

        metadata = self._ensure_metadata(conversation)
        attachments = metadata.get("attachments", [])
        incoming_ids = {item.get("file_id") for item in attachment_infos}
        attachments = [item for item in attachments if item.get("file_id") not in incoming_ids]
        attachments.extend(attachment_infos)
        metadata["attachments"] = attachments
        await self._save_metadata(conversation, metadata)
        return attachment_infos

    async def update_attachment_status(
        self, conversation_id: int, file_id: str, status: str, update_fields: dict | None = None
    ) -> dict | None:
        conversation = await self._get_conversation_by_id(conversation_id)
        if not conversation:
            return None

        metadata = self._ensure_metadata(conversation)
        attachments = metadata.get("attachments", [])
        target = None
        for item in attachments:
            if item.get("file_id") == file_id:
                item["status"] = status
                if update_fields:
                    item.update(update_fields)
                target = item
                break

        if target is not None:
            metadata["attachments"] = attachments
            await self._save_metadata(conversation, metadata)
        return target

    async def bind_attachments_to_request(
        self, conversation_id: int, request_id: str, file_ids: list[str]
    ) -> list[dict]:
        conversation = await self._get_conversation_by_id(conversation_id)
        if not conversation or not request_id or not file_ids:
            return []

        file_id_set = {str(file_id).strip() for file_id in file_ids if str(file_id).strip()}
        if not file_id_set:
            return []

        metadata = self._ensure_metadata(conversation)
        attachments = metadata.get("attachments", [])
        changed = False

        for item in attachments:
            if item.get("file_id") not in file_id_set:
                continue
            if item.get("request_id"):
                continue
            item["request_id"] = request_id
            changed = True

        if changed:
            metadata["attachments"] = attachments
            await self._save_metadata(conversation, metadata)
        return [dict(item) for item in attachments if item.get("request_id") == request_id]

    async def get_attachments_by_request_id(self, conversation_id: int, request_id: str) -> list[dict]:
        attachments = await self.get_attachments(conversation_id)
        return [item for item in attachments if item.get("request_id") == request_id]

    async def remove_attachment(self, conversation_id: int, file_id: str) -> bool:
        conversation = await self._get_conversation_by_id(conversation_id)
        if not conversation:
            return False

        metadata = self._ensure_metadata(conversation)
        attachments = metadata.get("attachments", [])
        new_attachments = [item for item in attachments if item.get("file_id") != file_id]

        if len(new_attachments) == len(attachments):
            return False

        metadata["attachments"] = new_attachments
        await self._save_metadata(conversation, metadata)
        return True
