from typing import TYPE_CHECKING
import structlog
from talemate.agents.base import (
    AgentAction,
    AgentActionConfig,
    AgentActionNote,
    AgentActionConditional,
)
import talemate.instance as instance
from talemate.util.dedupe import compile_text_to_sentences, compile_sentences_to_length

if TYPE_CHECKING:
    from talemate.tale_mate import Character

__all__ = ["MemoryRAGMixin"]

log = structlog.get_logger()

ai_assisted_condition = AgentActionConditional(
    attribute="use_long_term_memory.config.retrieval_method",
    value=["queries", "questions"],
)


class MemoryRAGMixin:
    @classmethod
    def add_actions(cls, actions: dict[str, AgentAction]):
        actions["use_long_term_memory"] = AgentAction(
            enabled=True,
            container=True,
            can_be_disabled=True,
            icon="mdi-brain",
            label="Long Term Memory",
            description="Will augment the context with long term memory based on similarity queries. Semantic similarity will ALWAYS be used, but you can configure an AI assisted method to use on top of it.",
            config={
                "retrieval_method": AgentActionConfig(
                    type="text",
                    label="AI Assistance",
                    description="AI assisted context retrieval method.",
                    value="direct",
                    choices=[
                        {
                            "label": "None (Semantic similarity only)",
                            "value": "direct",
                        },
                        {
                            "label": "Context queries generated by AI",
                            "value": "queries",
                        },
                        {
                            "label": "AI compiled question and answers (slow)",
                            "value": "questions",
                        },
                    ],
                    note_on_value={
                        "queries": AgentActionNote(
                            color="primary",
                            text="Will send one additional request to the AI to compile a set of similarity queries to run against the long term memory (+1 request)",
                        ),
                        "questions": AgentActionNote(
                            color="primary",
                            text="Will first prompt the LLM to compile a set of questions and then answer them (+2 requests)",
                        ),
                    },
                ),
                "num_messages": AgentActionConfig(
                    type="number",
                    label="Number of Messages",
                    description="When retrieving context based on semantic similarity, this is the number of messages to consider (going back from the most recent message)",
                    value=3,
                    min=1,
                    max=25,
                    step=1,
                ),
                "number_of_queries": AgentActionConfig(
                    type="number",
                    label="Number of Queries",
                    description="The number of queries to use when retrieving context from the long term memory.",
                    value=3,
                    min=1,
                    max=10,
                    step=1,
                    condition=ai_assisted_condition,
                ),
                "answer_length": AgentActionConfig(
                    type="text",
                    label="Answer Length",
                    description="The maximum length of long term memory response.",
                    value="512",
                    choices=[
                        {"label": "Short (256)", "value": "256"},
                        {"label": "Medium (512)", "value": "512"},
                        {"label": "Long (1024)", "value": "1024"},
                    ],
                    condition=ai_assisted_condition,
                ),
                "cache": AgentActionConfig(
                    type="bool",
                    label="Cache RAG results",
                    description="Cache the long term memory for faster retrieval.",
                    note="This is a cross-agent cache, assuming they use the same options.",
                    value=True,
                ),
            },
        )

    # config property helpers

    @property
    def long_term_memory_enabled(self):
        return self.actions["use_long_term_memory"].enabled

    @property
    def long_term_memory_retrieval_method(self):
        return self.actions["use_long_term_memory"].config["retrieval_method"].value

    @property
    def long_term_memory_number_of_queries(self):
        return self.actions["use_long_term_memory"].config["number_of_queries"].value

    @property
    def long_term_memory_answer_length(self):
        return int(self.actions["use_long_term_memory"].config["answer_length"].value)

    @property
    def long_term_memory_cache(self):
        return self.actions["use_long_term_memory"].config["cache"].value

    @property
    def long_term_memory_num_messages(self):
        return self.actions["use_long_term_memory"].config["num_messages"].value

    @property
    def long_term_memory_cache_key(self):
        """
        Build the key from the various options
        """

        parts = [
            self.long_term_memory_retrieval_method,
            self.long_term_memory_number_of_queries,
            self.long_term_memory_answer_length,
        ]

        return "-".join(map(str, parts))

    def connect(self, scene):
        super().connect(scene)

        # new scene, reset cache
        scene.rag_cache = {}

    # methods

    async def rag_set_cache(self, content: list[str]):
        self.scene.rag_cache[self.long_term_memory_cache_key] = {
            "content": content,
            "fingerprint": self.scene.history[-1].fingerprint
            if self.scene.history
            else 0,
        }

    async def rag_get_cache(self) -> list[str] | None:
        if not self.long_term_memory_cache:
            return None

        fingerprint = self.scene.history[-1].fingerprint if self.scene.history else 0
        cache = self.scene.rag_cache.get(self.long_term_memory_cache_key)

        if cache and cache["fingerprint"] == fingerprint:
            return cache["content"]

        return None

    async def rag_build(
        self,
        character: "Character | None" = None,
        prompt: str = "",
        sub_instruction: str = "",
    ) -> list[str]:
        """
        Builds long term memory to be inserted into a prompt
        """

        if not self.long_term_memory_enabled:
            return []

        cached = await self.rag_get_cache()

        if cached:
            log.debug(
                "Using cached long term memory",
                agent=self.agent_type,
                key=self.long_term_memory_cache_key,
            )
            return cached

        memory_context = ""
        semantic_context = await self.semantic_context(
            num_messages=self.long_term_memory_num_messages
        )
        retrieval_method = self.long_term_memory_retrieval_method

        if retrieval_method == "direct":
            # configuration is set to only use direct semantic matched context
            return semantic_context

        if not sub_instruction:
            if character:
                sub_instruction = f"continue the scene as {character.name}"
            elif hasattr(self, "rag_build_sub_instruction"):
                sub_instruction = await self.rag_build_sub_instruction()

        if not sub_instruction:
            sub_instruction = "continue the scene"

        world_state = instance.get_agent("world_state")

        if not prompt:
            prompt = self.scene.context_history(
                keep_director=False,
                budget=int(self.client.max_token_length * 0.75),
            )

        if isinstance(prompt, list):
            prompt = "\n".join(prompt)

        log.debug(
            "memory_rag_mixin.build_prompt_default_memory",
            direct=False,
            version=retrieval_method,
        )

        if retrieval_method == "questions":
            memory_context = (
                await world_state.analyze_text_and_extract_context(
                    prompt,
                    sub_instruction,
                    include_character_context=True,
                    response_length=self.long_term_memory_answer_length,
                    num_queries=self.long_term_memory_number_of_queries,
                    extra_context=semantic_context,
                )
            ).split("\n")
        elif retrieval_method == "queries":
            memory_context = (
                await world_state.analyze_text_and_extract_context_via_queries(
                    prompt,
                    sub_instruction,
                    include_character_context=True,
                    response_length=self.long_term_memory_answer_length,
                    num_queries=self.long_term_memory_number_of_queries,
                    extra_context=semantic_context,
                )
            )

        complete_context = list(set(semantic_context + memory_context))

        await self.rag_set_cache(complete_context)

        return complete_context

    async def semantic_context(
        self,
        num_messages: int = 3,
        min_query_length: int = 100,
        max_response_tokens: int = 1024,
    ):
        """Will retrieve context from the long term memory based on semantic similarity

        Args:
            num_messages (int, optional): The number of messages to consider (going back from the most recent message). Defaults to 3.
            min_query_length (int, optional): The minimum length of a query. Defaults to 100.
            max_response_tokens (int, optional): The maximum length of the response. Defaults to 1024.

        Returns:
            list[str]: The context retrieved from the long term memory
        """

        history = list(
            map(
                str,
                self.scene.collect_messages(
                    max_iterations=100,
                    max_messages=num_messages,
                    typ=["character", "narrator", "director"],
                ),
            )
        )
        log.debug(
            "memory_rag_mixin.build_prompt_default_memory",
            history=history,
            direct=True,
        )
        memory = instance.get_agent("memory")

        history_sentences = []
        for item in history:
            if not item.strip():
                continue
            sentences = compile_text_to_sentences(item)
            for sentence in sentences:
                if not sentence[1].strip():
                    continue
                history_sentences.append(sentence[1])

        history_sentences = compile_sentences_to_length(
            history_sentences, min_query_length
        )

        queries = [i for i in history + history_sentences if i.strip()]

        for query in queries:
            log.debug("memory_rag_mixin.build_prompt_default_memory", query=query)

        context = await memory.multi_query(
            queries, max_tokens=max_response_tokens, iterate=5
        )

        return context
