"""MessagesMixin — message building, memory, system prompts."""
from __future__ import annotations
from typing import Any

from logos_adk.agent._constants import _UNSET
from logos_adk.memory import MemoryEntry
from logos_adk.types import Message
from logos_adk.token_counter import count_text as _count_text


class MessagesMixin:
    async def _remember(
        self,
        messages: list[Message],
        state: Any,
        session_id: str | None,
    ) -> list[MemoryEntry]:
        """Collect memory entries relevant to the current turn."""
        entries: list[MemoryEntry] = []
        for memory in self._config.memories:
            entries.extend(await memory.remember(messages, state, session_id=session_id))
        return entries

    def _memory_messages(self, entries: list[MemoryEntry]) -> list[Message]:
        """Render retrieved memories into system messages, respecting token budget."""
        if not entries:
            return []

        budget = self._config.context_window.memory_max_tokens

        if budget <= 0:
            # No budget — include all entries
            return [
                Message(
                    role="system",
                    content=(
                        "Relevant memory:\n" + "\n".join(f"- {entry.content}" for entry in entries)
                    ),
                )
            ]

        # Include entries until we'd exceed the budget
        included: list[str] = []
        used_tokens = _count_text("Relevant memory:\n", model="")
        for entry in entries:
            line = f"- {entry.content}"
            line_tokens = _count_text(line, model="")
            if used_tokens + line_tokens > budget and included:
                break
            included.append(line)
            used_tokens += line_tokens

        if not included:
            return []

        return [
            Message(
                role="system",
                content="Relevant memory:\n" + "\n".join(included),
            )
        ]

    async def _commit_memories(
        self,
        messages: list[Message],
        response: Any,
        state: Any,
        session_id: str | None,
    ) -> None:
        """Persist memory after a completed agent turn."""
        for memory in self._config.memories:
            await memory.commit(messages, response, state, session_id=session_id)

    def _normalize_prompt_input(
        self,
        prompt: str | list,
    ) -> tuple[str, list]:
        if isinstance(prompt, str):
            return prompt, []
        from logos_adk.types import normalize_content_parts, content_text
        parts = normalize_content_parts(prompt)
        return content_text("", parts), parts

    def _resolved_system_prompt(self, system_prompt_override: Any = _UNSET) -> str | None:
        # If override is provided, skip cache
        if system_prompt_override is not _UNSET:
            return self._compute_system_prompt(system_prompt_override)
        # Check run-scoped cache
        if self._run_cache_depth > 0 and self._cached_system_prompt is not _UNSET:
            return self._cached_system_prompt  # type: ignore[return-value]
        result = self._compute_system_prompt(_UNSET)
        if self._run_cache_depth > 0:
            self._cached_system_prompt = result
        return result

    def _compute_system_prompt(self, system_prompt_override: Any) -> str | None:
        base_prompt = self._config.system_prompt if system_prompt_override is _UNSET else system_prompt_override
        role_prompt = self._role_system_prompt()

        # --- Self-Learning Injection ---
        lessons = self._get_learned_lessons()
        learning_prompt = None
        if lessons:
            learning_prompt = "### LEARNED FROM PREVIOUS FEEDBACK:\n" + "\n".join(
                f"- {lesson}" for lesson in lessons
            )
        # -------------------------------
        budget_prompt = self._budget_system_prompt()

        parts = [
            part.strip()
            for part in [role_prompt, base_prompt, learning_prompt, budget_prompt]
            if isinstance(part, str) and part.strip()
        ]
        if not parts:
            return None
        return "\n\n".join(parts)

    def _get_learned_lessons(self) -> list[str]:
        """Retrieve learned lessons from agent metadata."""
        lessons = self._config.metadata.get("learned_lessons", [])
        if not isinstance(lessons, list):
            return []
        return lessons

    def learn_from_feedback(self, feedback: str) -> None:
        """Save a new lesson learned from user feedback."""
        if not feedback or len(feedback) < 3:
            return
        lessons = self._config.metadata.setdefault("learned_lessons", [])
        if feedback not in lessons:
            lessons.append(feedback)
            # Persist if needed
            if self._config_path:
                try:
                    self.save_config(self._config_path)
                except Exception:
                    pass

    def _role_system_prompt(self) -> str | None:
        profile = self.crew_profile
        if not profile:
            return None
        lines = ["Crew Profile:"]
        if "role" in profile:
            lines.append(f"Role: {profile['role']}")
        if "goal" in profile:
            lines.append(f"Goal: {profile['goal']}")
        if "backstory" in profile:
            lines.append(f"Backstory: {profile['backstory']}")
        return "\n".join(lines)
