"""Adaptive learning helpers for agent memory systems."""
from __future__ import annotations

from datetime import UTC, datetime
from typing import Any

from logos_adk.memory import MemoryBackend, MemoryEntry


class MemoryLearningEngine:
    """Adaptive scoring, pruning, and summary tuning for memory entries."""

    def __init__(
        self,
        *,
        archive_threshold: float = 0.25,
        delete_threshold: float = 0.1,
        max_active_entries: int = 50,
        archive_after_days: int = 30,
        max_context_limit: int = 10,
    ) -> None:
        self.archive_threshold = max(0.0, min(1.0, archive_threshold))
        self.delete_threshold = max(0.0, min(self.archive_threshold, delete_threshold))
        self.max_active_entries = max(5, max_active_entries)
        self.archive_after_days = max(1, archive_after_days)
        self.max_context_limit = max(2, max_context_limit)

    def initialize_entry(self, entry: MemoryEntry, *, kind: str = "conversation_turn") -> MemoryEntry:
        entry.metadata.setdefault("kind", kind)
        entry.metadata.setdefault("usefulness_score", 0.5)
        entry.metadata.setdefault("trust_score", 0.55)
        entry.metadata.setdefault("helpful_count", 0)
        entry.metadata.setdefault("harmful_count", 0)
        entry.metadata.setdefault("access_count", 0)
        entry.metadata.setdefault("last_accessed_at", entry.created_at.isoformat())
        return entry

    async def observe_remembered(
        self,
        backend: MemoryBackend,
        entries: list[MemoryEntry],
    ) -> None:
        for entry in entries:
            entry.mark_accessed()
            access_boost = min(0.03, 0.005 * max(1, entry.access_count))
            entry.apply_learning_feedback(
                helpful=True,
                trust_delta=0.0,
                usefulness_delta=access_boost,
                reason="memory_recalled",
            )
            await backend.store(entry)

    def rerank_entries(
        self,
        entries: list[MemoryEntry],
        *,
        lexical_scores: dict[str, float],
    ) -> list[MemoryEntry]:
        ranked = sorted(
            (entry for entry in entries if not entry.is_archived),
            key=lambda entry: (
                lexical_scores.get(entry.id, 0.0) * 0.55
                + entry.utility_score * 0.35
                + entry.recency_score() * 0.10
            ),
            reverse=True,
        )
        return ranked

    def resolve_context_limit(self, entries: list[MemoryEntry], *, base_limit: int) -> int:
        if not entries:
            return max(1, base_limit)
        average_utility = sum(entry.utility_score for entry in entries[:base_limit]) / min(len(entries), base_limit)
        if average_utility >= 0.8:
            return max(1, base_limit - 1)
        if average_utility <= 0.35:
            return min(self.max_context_limit, base_limit + 2)
        if average_utility <= 0.5:
            return min(self.max_context_limit, base_limit + 1)
        return base_limit

    def summary_settings(self, *, historical_count: int, token_pressure: bool = False) -> dict[str, Any]:
        if historical_count >= 24 or token_pressure:
            return {"strategy": "compact", "line_limit": 48, "max_lines": 10}
        if historical_count >= 12:
            return {"strategy": "balanced", "line_limit": 80, "max_lines": 14}
        return {"strategy": "detailed", "line_limit": 120, "max_lines": 18}

    async def apply_feedback(
        self,
        backend: MemoryBackend,
        entries: list[MemoryEntry],
        *,
        helpful: bool,
        reason: str,
    ) -> None:
        trust_delta = 0.08 if helpful else -0.14
        usefulness_delta = 0.1 if helpful else -0.18
        for entry in entries:
            entry.apply_learning_feedback(
                helpful=helpful,
                trust_delta=trust_delta,
                usefulness_delta=usefulness_delta,
                reason=reason,
            )
            if not helpful and entry.utility_score <= self.archive_threshold:
                entry.archive(reason="negative_feedback")
            await backend.store(entry)

    async def sync_feedback(
        self,
        backend: MemoryBackend,
        *,
        entry_ids: list[str],
        session_id: str | None,
        helpful: bool,
        reason: str,
    ) -> list[MemoryEntry]:
        unique_ids = [entry_id for entry_id in dict.fromkeys(entry_ids) if entry_id]
        if not unique_ids:
            return []
        candidates = await backend.search(
            "",
            limit=max(100, self.max_active_entries * 5, len(unique_ids) * 10),
            session_id=session_id,
        )
        selected = [entry for entry in candidates if entry.id in set(unique_ids)]
        if not selected:
            return []
        await self.apply_feedback(backend, selected, helpful=helpful, reason=reason)
        return selected

    async def enforce_forgetting_policy(
        self,
        backend: MemoryBackend,
        *,
        session_id: str | None,
    ) -> None:
        entries = await backend.search("", limit=max(self.max_active_entries * 3, 60), session_id=session_id)
        if not entries:
            return
        now = datetime.now(UTC)
        active_entries = [entry for entry in entries if not entry.is_archived]
        archived_entries = [entry for entry in entries if entry.is_archived]

        ranked_active = sorted(active_entries, key=lambda entry: (entry.utility_score, entry.created_at))
        overflow = max(0, len(active_entries) - self.max_active_entries)
        for entry in ranked_active[:overflow]:
            entry.archive(reason="context_budget_prune")
            await backend.store(entry)

        for entry in active_entries:
            if entry.utility_score > self.archive_threshold:
                continue
            age_days = max(0.0, (now - entry.created_at).total_seconds() / 86400.0)
            if age_days >= self.archive_after_days:
                entry.archive(reason="low_value_retention")
                await backend.store(entry)

        for entry in archived_entries:
            archived_at = entry.archived_at or entry.created_at
            archived_days = max(0.0, (now - archived_at).total_seconds() / 86400.0)
            if archived_days >= self.archive_after_days and entry.utility_score <= self.delete_threshold:
                await backend.delete(entry.id)


__all__ = ["MemoryLearningEngine"]
