"""Memory system for Logos ADK."""
from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime, UTC
from typing import Any
from uuid import uuid4

from logos_adk.errors import ConfigError
from logos_adk.memory.retrieval import (
    VALID_RETRIEVALS,
    retrieval_queries,
    score_entry,
)
from logos_adk.state import AgentState
from logos_adk.types import Message, Response, content_text


@dataclass
class MemoryEntry:
    """A stored unit of memory."""

    content: str
    id: str = field(default_factory=lambda: str(uuid4()))
    embedding: list[float] | None = None
    metadata: dict[str, Any] = field(default_factory=dict)
    created_at: datetime = field(default_factory=lambda: datetime.now(UTC))
    session_id: str | None = None

    def _get_float(self, key: str, default: float) -> float:
        value = self.metadata.get(key, default)
        try:
            return float(value)
        except (TypeError, ValueError):
            return default

    def _get_int(self, key: str, default: int) -> int:
        value = self.metadata.get(key, default)
        try:
            return int(value)
        except (TypeError, ValueError):
            return default

    @property
    def usefulness_score(self) -> float:
        return max(0.0, min(1.0, self._get_float("usefulness_score", 0.5)))

    @property
    def trust_score(self) -> float:
        return max(0.0, min(1.0, self._get_float("trust_score", 0.5)))

    @property
    def helpful_count(self) -> int:
        return max(0, self._get_int("helpful_count", 0))

    @property
    def harmful_count(self) -> int:
        return max(0, self._get_int("harmful_count", 0))

    @property
    def access_count(self) -> int:
        return max(0, self._get_int("access_count", 0))

    @property
    def last_accessed_at(self) -> datetime | None:
        raw = self.metadata.get("last_accessed_at")
        if not raw:
            return None
        if isinstance(raw, datetime):
            return raw
        try:
            return datetime.fromisoformat(str(raw))
        except ValueError:
            return None

    @property
    def archived_at(self) -> datetime | None:
        raw = self.metadata.get("archived_at")
        if not raw:
            return None
        if isinstance(raw, datetime):
            return raw
        try:
            return datetime.fromisoformat(str(raw))
        except ValueError:
            return None

    @property
    def is_archived(self) -> bool:
        return self.archived_at is not None

    def recency_score(self, *, now: datetime | None = None) -> float:
        baseline = self.last_accessed_at or self.created_at
        current = now or datetime.now(UTC)
        age_seconds = max(0.0, (current - baseline).total_seconds())
        if age_seconds <= 3600:
            return 1.0
        if age_seconds <= 86400:
            return 0.8
        if age_seconds <= 7 * 86400:
            return 0.55
        if age_seconds <= 30 * 86400:
            return 0.35
        return 0.15

    @property
    def utility_score(self) -> float:
        return max(
            0.0,
            min(
                1.0,
                (self.usefulness_score * 0.5)
                + (self.trust_score * 0.25)
                + (self.recency_score() * 0.25),
            ),
        )

    def mark_accessed(self, when: datetime | None = None) -> None:
        current = when or datetime.now(UTC)
        self.metadata["last_accessed_at"] = current.isoformat()
        self.metadata["access_count"] = self.access_count + 1

    def apply_learning_feedback(
        self,
        *,
        helpful: bool,
        trust_delta: float = 0.0,
        usefulness_delta: float = 0.0,
        reason: str | None = None,
    ) -> None:
        if helpful:
            self.metadata["helpful_count"] = self.helpful_count + 1
        else:
            self.metadata["harmful_count"] = self.harmful_count + 1
        next_usefulness = max(0.0, min(1.0, self.usefulness_score + usefulness_delta))
        next_trust = max(0.0, min(1.0, self.trust_score + trust_delta))
        self.metadata["usefulness_score"] = round(next_usefulness, 4)
        self.metadata["trust_score"] = round(next_trust, 4)
        if reason:
            self.metadata["last_feedback_reason"] = reason

    def archive(self, *, reason: str) -> None:
        self.metadata["archived_at"] = datetime.now(UTC).isoformat()
        self.metadata["archived_reason"] = reason

    def restore(self) -> None:
        self.metadata.pop("archived_at", None)
        self.metadata.pop("archived_reason", None)


class MemoryBackend(ABC):
    """Storage backend for memory entries."""

    @abstractmethod
    async def store(self, entry: MemoryEntry) -> None:
        ...

    @abstractmethod
    async def search(
        self,
        query: str,
        limit: int = 10,
        session_id: str | None = None,
    ) -> list[MemoryEntry]:
        ...

    @abstractmethod
    async def delete(self, entry_id: str) -> None:
        ...


class Memory(ABC):
    """Strategy for remembering and committing memory."""

    @abstractmethod
    async def remember(
        self,
        messages: list[Message],
        state: AgentState,
        session_id: str | None = None,
    ) -> list[MemoryEntry]:
        ...

    @abstractmethod
    async def commit(
        self,
        messages: list[Message],
        response: Response,
        state: AgentState,
        session_id: str | None = None,
    ) -> None:
        ...


class SessionMemory(Memory):
    """Short-term memory that manages the in-session message window."""

    _VALID_STRATEGIES = {"sliding_window", "summarize", "token_limit"}

    def __init__(
        self,
        *,
        strategy: str = "sliding_window",
        max_messages: int = 20,
        max_tokens: int | None = None,
        learning: Any | None = None,
    ) -> None:
        if strategy not in self._VALID_STRATEGIES:
            valid = ", ".join(sorted(self._VALID_STRATEGIES))
            raise ConfigError(
                f"Unsupported session memory strategy: {strategy}. Valid strategies: {valid}"
            )
        if not isinstance(max_messages, int):
            raise ConfigError("SessionMemory.max_messages must be an integer")
        if max_messages < 0:
            raise ConfigError("SessionMemory.max_messages must be >= 0")
        if max_tokens is not None and not isinstance(max_tokens, int):
            raise ConfigError("SessionMemory.max_tokens must be an integer or None")
        if max_tokens is not None and max_tokens < 0:
            raise ConfigError("SessionMemory.max_tokens must be >= 0")
        if strategy == "token_limit" and max_tokens is None:
            raise ConfigError(
                "SessionMemory(strategy='token_limit') requires max_tokens to be set"
            )

        self.strategy = strategy
        self.max_messages = max_messages
        self.max_tokens = max_tokens
        self.learning = learning

    def _clip_text(self, text: str, limit: int = 120) -> str:
        text = " ".join(text.split())
        if len(text) <= limit:
            return text
        return f"{text[: limit - 3].rstrip()}..."

    def _estimate_text_tokens(self, text: str) -> int:
        text = text.strip()
        if not text:
            return 0
        return len(text.split())

    def _estimate_message_tokens(self, message: Message) -> int:
        total = 1 + self._estimate_text_tokens(content_text(message.content, message.content_parts))

        if message.tool_calls:
            for tool_call in message.tool_calls:
                total += 1 + self._estimate_text_tokens(tool_call.name)
                total += self._estimate_text_tokens(str(tool_call.arguments))

        if message.tool_result is not None:
            total += 1
            total += self._estimate_text_tokens(
                content_text(
                    str(message.tool_result.content) if message.tool_result.content is not None else "",
                    message.tool_result.content_parts,
                )
            )
            total += self._estimate_text_tokens(message.tool_result.error or "")

        return total

    def _estimate_messages_tokens(self, messages: list[Message]) -> int:
        return sum(self._estimate_message_tokens(message) for message in messages)

    def _apply_token_limit(self, messages: list[Message]) -> list[Message]:
        if self.max_tokens is None:
            return messages[-self.max_messages :]

        if self.max_tokens <= 0:
            return []

        recent_messages = list(messages)
        historical_messages: list[Message] = []

        while (
            len(recent_messages) > 1
            and self._estimate_messages_tokens(recent_messages) > self.max_tokens
        ):
            historical_messages.append(recent_messages.pop(0))

        if not historical_messages:
            return recent_messages

        summary_settings = {"strategy": "compact", "line_limit": 48, "max_lines": 10}
        if self.learning is not None:
            summary_settings = self.learning.summary_settings(
                historical_count=len(historical_messages),
                token_pressure=True,
            )
        summary_message = self._build_summary_message(
            historical_messages,
            compact=summary_settings["strategy"] == "compact",
            line_limit=summary_settings["line_limit"],
            max_lines=summary_settings["max_lines"],
        )
        if summary_message is None:
            return recent_messages

        candidate_messages = [summary_message] + recent_messages
        if self._estimate_messages_tokens(candidate_messages) <= self.max_tokens:
            return candidate_messages

        return recent_messages

    def _summary_line(self, message: Message, *, limit: int = 120) -> str | None:
        if message.role == "tool":
            if message.tool_result is not None:
                if message.tool_result.error:
                    return (
                        f"Tool {message.tool_result.call_id} failed: "
                        f"{self._clip_text(message.tool_result.error, limit=min(limit, 80))}"
                    )
                tool_text = content_text(
                    str(message.tool_result.content) if message.tool_result.content is not None else "",
                    message.tool_result.content_parts,
                )
                return (
                    f"Tool {message.tool_result.call_id} returned "
                    f"{self._clip_text(tool_text, limit=min(limit, 80))}"
                )
            tool_message_text = content_text(message.content, message.content_parts)
            if tool_message_text:
                return f"Tool output: {self._clip_text(tool_message_text, limit=min(limit, 80))}"
            return None

        if message.role == "assistant" and message.tool_calls:
            tool_names = ", ".join(tool_call.name for tool_call in message.tool_calls)
            assistant_text = content_text(message.content, message.content_parts)
            if assistant_text:
                return (
                    f"Assistant decided to call tools ({tool_names}) after saying "
                    f"{self._clip_text(assistant_text, limit=limit)}"
                )
            return f"Assistant decided to call tools: {tool_names}"

        message_text = content_text(message.content, message.content_parts)
        if not message_text:
            return None

        role_label = {
            "user": "User",
            "assistant": "Assistant",
            "system": "System",
        }.get(message.role, message.role.capitalize())
        return f"{role_label}: {self._clip_text(message_text, limit=limit)}"

    def _build_summary_message(
        self,
        messages: list[Message],
        *,
        compact: bool = False,
        line_limit: int | None = None,
        max_lines: int | None = None,
    ) -> Message | None:
        summary_lines: list[str] = []
        effective_line_limit = line_limit or (48 if compact else 120)

        for message in messages:
            message_text = content_text(message.content, message.content_parts)
            if message.role == "system" and message_text.startswith("Session summary:\n"):
                existing = message_text.removeprefix("Session summary:\n").strip()
                if existing:
                    summary_lines.extend(
                        line for line in existing.splitlines() if line.strip()
                    )
                continue

            summary_line = self._summary_line(message, limit=effective_line_limit)
            if summary_line is not None:
                summary_lines.append(f"- {summary_line}")

        if not summary_lines:
            return None
        if max_lines is not None and max_lines > 0:
            summary_lines = summary_lines[:max_lines]

        return Message(
            role="system",
            content="Session summary:\n" + "\n".join(summary_lines),
        )

    async def remember(
        self,
        messages: list[Message],
        state: AgentState,
        session_id: str | None = None,
    ) -> list[MemoryEntry]:
        return []

    async def commit(
        self,
        messages: list[Message],
        response: Response,
        state: AgentState,
        session_id: str | None = None,
    ) -> None:
        if self.max_messages <= 0:
            state.messages = []
            return

        if self.strategy == "summarize":
            recent_messages = state.messages[-self.max_messages :]
            historical_messages = state.messages[: -self.max_messages]
            summary_settings = {"strategy": "detailed", "line_limit": 120, "max_lines": 18}
            if self.learning is not None:
                summary_settings = self.learning.summary_settings(
                    historical_count=len(historical_messages),
                    token_pressure=False,
                )
                state.metadata["adk_memory_summary_strategy"] = summary_settings["strategy"]
            summary_message = self._build_summary_message(
                historical_messages,
                compact=summary_settings["strategy"] == "compact",
                line_limit=summary_settings["line_limit"],
                max_lines=summary_settings["max_lines"],
            )
            state.messages = (
                ([summary_message] if summary_message is not None else []) + recent_messages
            )
            return

        if self.strategy == "token_limit":
            if self.learning is not None:
                summary_settings = self.learning.summary_settings(
                    historical_count=max(0, len(state.messages) - self.max_messages),
                    token_pressure=True,
                )
                state.metadata["adk_memory_summary_strategy"] = summary_settings["strategy"]
            state.messages = self._apply_token_limit(state.messages)
            return

        state.messages = state.messages[-self.max_messages :]


class LongTermMemory(Memory):
    """Persistent memory backed by a MemoryBackend."""

    _VALID_RETRIEVALS = VALID_RETRIEVALS

    def __init__(
        self,
        *,
        backend: MemoryBackend,
        retrieval: str = "keyword",
        limit: int = 5,
        embed_model: str | None = None,
        learning: Any | None = None,
    ) -> None:
        if retrieval not in self._VALID_RETRIEVALS:
            valid = ", ".join(sorted(self._VALID_RETRIEVALS))
            raise ConfigError(
                f"Unsupported long-term memory retrieval: {retrieval}. "
                f"Valid retrieval modes: {valid}"
            )
        if not isinstance(limit, int):
            raise ConfigError("LongTermMemory.limit must be an integer")
        if limit <= 0:
            raise ConfigError("LongTermMemory.limit must be > 0")

        self.backend = backend
        self.retrieval = retrieval
        self.limit = limit
        self.embed_model = embed_model
        self.learning = learning

    def _scope_session_id(self, session_id: str | None) -> str | None:
        return session_id

    def _last_user_query(self, messages: list[Message]) -> str:
        for message in reversed(messages):
            message_text = content_text(message.content, message.content_parts)
            if message.role == "user" and message_text:
                return message_text
        return ""

    async def remember(
        self,
        messages: list[Message],
        state: AgentState,
        session_id: str | None = None,
    ) -> list[MemoryEntry]:
        query = self._last_user_query(messages)
        if not query:
            return []

        scoped_session_id = self._scope_session_id(session_id)
        candidates: dict[str, MemoryEntry] = {}
        lexical_scores: dict[str, float] = {}

        for search_query in retrieval_queries(query, self.retrieval):
            results = await self.backend.search(
                search_query,
                limit=max(self.limit * 2, self.limit),
                session_id=scoped_session_id,
            )
            for entry in results:
                if entry.is_archived:
                    continue
                candidates[entry.id] = entry
                lexical_scores[entry.id] = max(
                    lexical_scores.get(entry.id, 0.0),
                    score_entry(query, entry, self.retrieval),
                )

        ranked_entries = sorted(
            candidates.values(),
            key=lambda entry: (lexical_scores.get(entry.id, 0.0), entry.created_at),
            reverse=True,
        )
        if self.learning is not None:
            ranked_entries = self.learning.rerank_entries(
                ranked_entries,
                lexical_scores=lexical_scores,
            )
            resolved_limit = self.learning.resolve_context_limit(
                ranked_entries,
                base_limit=self.limit,
            )
            state.metadata["adk_memory_context_limit"] = resolved_limit
            state.metadata["adk_memory_entry_ids"] = [entry.id for entry in ranked_entries[:resolved_limit]]
            await self.learning.observe_remembered(self.backend, ranked_entries[:resolved_limit])
            return ranked_entries[:resolved_limit]
        state.metadata["adk_memory_context_limit"] = self.limit
        state.metadata["adk_memory_entry_ids"] = [entry.id for entry in ranked_entries[: self.limit]]
        return ranked_entries[: self.limit]

    async def commit(
        self,
        messages: list[Message],
        response: Response,
        state: AgentState,
        session_id: str | None = None,
    ) -> None:
        user_query = self._last_user_query(messages)
        response_text = content_text(response.content, response.content_parts)
        if not user_query and not response_text:
            return

        entry = MemoryEntry(
            content=f"User: {user_query}\nAssistant: {response_text}".strip(),
            metadata={
                "kind": "conversation_turn",
                "retrieval": self.retrieval,
            },
            session_id=self._scope_session_id(session_id),
        )
        if self.learning is not None:
            entry = self.learning.initialize_entry(entry)
        await self.backend.store(entry)
        if self.learning is not None:
            await self.learning.enforce_forgetting_policy(
                self.backend,
                session_id=self._scope_session_id(session_id),
            )


class SharedMemory(LongTermMemory):
    """Long-term memory shared across sessions and agents."""

    def _scope_session_id(self, session_id: str | None) -> str | None:
        return None


__all__ = [
    "MemoryEntry",
    "MemoryBackend",
    "Memory",
    "SessionMemory",
    "LongTermMemory",
    "SharedMemory",
]
