from collections import defaultdict
from collections.abc import Sequence

from sqlalchemy import select
from sqlalchemy.orm import sessionmaker

from core.app.app_config.features.file_upload.manager import FileUploadConfigManager
from core.app.file_access import DatabaseFileAccessController
from core.model_manager import ModelInstance
from core.prompt.utils.extract_thread_messages import extract_thread_messages
from extensions.ext_database import db
from factories import file_factory
from graphon.file import file_manager
from graphon.model_runtime.entities import (
    AssistantPromptMessage,
    ImagePromptMessageContent,
    PromptMessage,
    PromptMessageRole,
    TextPromptMessageContent,
    UserPromptMessage,
)
from graphon.model_runtime.entities.message_entities import PromptMessageContentUnionTypes
from models.model import AppMode, Conversation, Message, MessageFile
from models.workflow import Workflow
from repositories.api_workflow_run_repository import APIWorkflowRunRepository
from repositories.factory import DifyAPIRepositoryFactory

_file_access_controller = DatabaseFileAccessController()


class TokenBufferMemory:
    def __init__(
        self,
        conversation: Conversation,
        model_instance: ModelInstance,
    ):
        self.conversation = conversation
        self.model_instance = model_instance
        self._workflow_run_repo: APIWorkflowRunRepository | None = None

    @property
    def workflow_run_repo(self) -> APIWorkflowRunRepository:
        if self._workflow_run_repo is None:
            session_maker = sessionmaker(bind=db.engine, expire_on_commit=False)
            self._workflow_run_repo = DifyAPIRepositoryFactory.create_api_workflow_run_repository(session_maker)
        return self._workflow_run_repo

    def _build_prompt_message_with_files(
        self,
        message_files: Sequence[MessageFile],
        text_content: str,
        message: Message,
        app_record,
        is_user_message: bool,
    ) -> PromptMessage:
        """
        Build prompt message with files.
        :param message_files: Sequence of MessageFile objects
        :param text_content: text content of the message
        :param message: Message object
        :param app_record: app record
        :param is_user_message: whether this is a user message
        :return: PromptMessage
        """
        match self.conversation.mode:
            case AppMode.AGENT_CHAT | AppMode.COMPLETION | AppMode.CHAT:
                file_extra_config = FileUploadConfigManager.convert(self.conversation.model_config)
            case AppMode.ADVANCED_CHAT | AppMode.WORKFLOW:
                app = self.conversation.app
                if not app:
                    raise ValueError("App not found for conversation")

                if not message.workflow_run_id:
                    raise ValueError("Workflow run ID not found")

                workflow_run = self.workflow_run_repo.get_workflow_run_by_id(
                    tenant_id=app.tenant_id, app_id=app.id, run_id=message.workflow_run_id
                )
                if not workflow_run:
                    raise ValueError(f"Workflow run not found: {message.workflow_run_id}")
                workflow = db.session.scalar(select(Workflow).where(Workflow.id == workflow_run.workflow_id))
                if not workflow:
                    raise ValueError(f"Workflow not found: {workflow_run.workflow_id}")
                file_extra_config = FileUploadConfigManager.convert(workflow.features_dict, is_vision=False)
            case _:
                raise AssertionError(f"Invalid app mode: {self.conversation.mode}")

        detail = ImagePromptMessageContent.DETAIL.HIGH
        if file_extra_config and app_record:
            file_objs = [
                file_factory.build_from_message_file(
                    message_file=message_file,
                    tenant_id=app_record.tenant_id,
                    access_controller=_file_access_controller,
                )
                for message_file in message_files
            ]
            if file_extra_config.image_config and file_extra_config.image_config.detail:
                detail = file_extra_config.image_config.detail
        else:
            file_objs = []

        if not file_objs:
            if is_user_message:
                return UserPromptMessage(content=text_content)
            else:
                return AssistantPromptMessage(content=text_content)
        else:
            prompt_message_contents: list[PromptMessageContentUnionTypes] = []
            for file in file_objs:
                prompt_message = file_manager.to_prompt_message_content(
                    file,
                    image_detail_config=detail,
                )
                prompt_message_contents.append(prompt_message)
            prompt_message_contents.append(TextPromptMessageContent(data=text_content))

            if is_user_message:
                return UserPromptMessage(content=prompt_message_contents)
            else:
                return AssistantPromptMessage(content=prompt_message_contents)

    def get_history_prompt_messages(
        self, max_token_limit: int = 2000, message_limit: int | None = None
    ) -> Sequence[PromptMessage]:
        """
        Get history prompt messages.
        :param max_token_limit: max token limit
        :param message_limit: message limit
        """
        app_record = self.conversation.app

        # fetch limited messages, and return reversed
        stmt = (
            select(Message).where(Message.conversation_id == self.conversation.id).order_by(Message.created_at.desc())
        )

        if message_limit and message_limit > 0:
            message_limit = min(message_limit, 500)
        else:
            message_limit = 500

        msg_limit_stmt = stmt.limit(message_limit)

        messages = db.session.scalars(msg_limit_stmt).all()

        # instead of all messages from the conversation, we only need to extract messages
        # that belong to the thread of last message
        thread_messages = extract_thread_messages(messages)

        # for newly created message, its answer is temporarily empty, we don't need to add it to memory
        if thread_messages and not thread_messages[0].answer and thread_messages[0].answer_tokens == 0:
            thread_messages.pop(0)

        messages = list(reversed(thread_messages))

        # Batch-load message files for the whole thread to avoid an N+1 query pattern.
        # Previously each message issued two MessageFile queries (user + assistant),
        # i.e. 2N+1 round-trips for N messages. We now use two batched queries keyed by
        # message_id, preserving the exact filter semantics (user files include rows
        # whose belongs_to is NULL).
        message_ids = [message.id for message in messages]
        user_files_by_message: dict[str, list[MessageFile]] = defaultdict(list)
        assistant_files_by_message: dict[str, list[MessageFile]] = defaultdict(list)
        if message_ids:
            for message_file in db.session.scalars(
                select(MessageFile).where(
                    MessageFile.message_id.in_(message_ids),
                    (MessageFile.belongs_to == "user") | (MessageFile.belongs_to.is_(None)),
                )
            ).all():
                user_files_by_message[message_file.message_id].append(message_file)

            for message_file in db.session.scalars(
                select(MessageFile).where(
                    MessageFile.message_id.in_(message_ids),
                    MessageFile.belongs_to == "assistant",
                )
            ).all():
                assistant_files_by_message[message_file.message_id].append(message_file)

        curr_message_tokens = 0
        prompt_messages: list[PromptMessage] = []
        for message in messages:
            # Process user message with files
            user_files = user_files_by_message.get(message.id, [])

            if user_files:
                user_prompt_message = self._build_prompt_message_with_files(
                    message_files=user_files,
                    text_content=message.query,
                    message=message,
                    app_record=app_record,
                    is_user_message=True,
                )
                prompt_messages.append(user_prompt_message)
            else:
                prompt_messages.append(UserPromptMessage(content=message.query))

            # Process assistant message with files
            assistant_files = assistant_files_by_message.get(message.id, [])

            if assistant_files:
                assistant_prompt_message = self._build_prompt_message_with_files(
                    message_files=assistant_files,
                    text_content=message.answer,
                    message=message,
                    app_record=app_record,
                    is_user_message=False,
                )
                prompt_messages.append(assistant_prompt_message)
            else:
                prompt_messages.append(AssistantPromptMessage(content=message.answer))

        if not prompt_messages:
            return []

        # prune the chat message if it exceeds the max token limit
        curr_message_tokens = self.model_instance.get_llm_num_tokens(prompt_messages)

        if curr_message_tokens > max_token_limit:
            while curr_message_tokens > max_token_limit and len(prompt_messages) > 1:
                prompt_messages.pop(0)
                curr_message_tokens = self.model_instance.get_llm_num_tokens(prompt_messages)

        return prompt_messages

    def get_history_prompt_text(
        self,
        human_prefix: str = "Human",
        ai_prefix: str = "Assistant",
        max_token_limit: int = 2000,
        message_limit: int | None = None,
    ) -> str:
        """
        Get history prompt text.
        :param human_prefix: human prefix
        :param ai_prefix: ai prefix
        :param max_token_limit: max token limit
        :param message_limit: message limit
        :return:
        """
        prompt_messages = self.get_history_prompt_messages(max_token_limit=max_token_limit, message_limit=message_limit)

        string_messages = []
        for m in prompt_messages:
            if m.role == PromptMessageRole.USER:
                role = human_prefix
            elif m.role == PromptMessageRole.ASSISTANT:
                role = ai_prefix
            else:
                continue

            if isinstance(m.content, list):
                inner_msg = ""
                for content in m.content:
                    match content:
                        case TextPromptMessageContent():
                            inner_msg += f"{content.data}\n"
                        case ImagePromptMessageContent():
                            inner_msg += "[image]\n"

                string_messages.append(f"{role}: {inner_msg.strip()}")
            else:
                message = f"{role}: {m.content}"
                string_messages.append(message)

        return "\n".join(string_messages)
