# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Agentic Memory Management

Proactively extracts facts from conversations and manages memory
lifecycle. Runs asynchronously after the agent response is sent,
so it adds zero latency to the user experience.

Two main functions:

1. **Memory Extraction** — Analyzes the conversation transcript
   and extracts structured facts (user preferences, org info,
   deadlines, corrections). Stores them via Memory.remember().

2. **Memory Decay** — Periodically reduces importance of
   unaccessed memories. Old low-importance memories are archived.

Usage:
    extractor = MemoryExtractor(provider, memory)

    # After a conversation completes:
    extractor.extract_async(user_message, assistant_response, channel)

    # Periodic maintenance (e.g., daily cron):
    decay_memories(memory, days_idle=30, decay_amount=1)

Integration:
    Called from Agent.chat() after final_text is returned.
    Uses a background thread so the user never waits.
"""

import json
import logging
import re
import threading
import time
from dataclasses import dataclass
from datetime import datetime
from typing import TYPE_CHECKING, Any, Dict, List, Optional

if TYPE_CHECKING:
    from .memory import Memory
    from .providers import LLMProvider

logger = logging.getLogger(__name__)


# ════════════════════════════════════════════════════════════════
# Extraction Data Structures
# ════════════════════════════════════════════════════════════════


@dataclass
class ExtractedMemory:
    """A single fact extracted from a conversation."""

    key: str
    value: str
    category: str  # user_info, org_info, deadline, preference, correction
    importance: int = 5  # 1-10
    supersedes: Optional[str] = None  # Key of memory this replaces

    def to_dict(self) -> dict:
        d = {
            "key": self.key,
            "value": self.value,
            "category": self.category,
            "importance": self.importance,
        }
        if self.supersedes:
            d["supersedes"] = self.supersedes
        return d

    @classmethod
    def from_dict(cls, data: dict) -> "ExtractedMemory":
        return cls(
            key=data.get("key", ""),
            value=data.get("value", ""),
            category=data.get("category", "fact"),
            importance=data.get("importance", 5),
            supersedes=data.get("supersedes"),
        )


# ════════════════════════════════════════════════════════════════
# Extraction Prompt
# ════════════════════════════════════════════════════════════════

EXTRACTION_SYSTEM = """You are a memory extraction system. Analyze the conversation
and extract important facts worth remembering for future interactions.

Respond with ONLY a JSON array (no markdown, no explanation):
[
  {
    "key": "short_identifier",
    "value": "the fact in a clear sentence",
    "category": "user_info|org_info|deadline|preference|correction",
    "importance": 5,
    "supersedes": "key_of_memory_this_replaces_or_null"
  }
]

Categories:
- user_info: Personal facts about the user (name, role, location, background)
- org_info: Organizational facts (board members, fiscal year, team structure)
- deadline: Time-sensitive items (grant due dates, meeting dates, renewals)
- preference: User preferences (communication style, tools, formats)
- correction: Updates to previously stored facts

Rules:
- Extract 0-5 facts per conversation (fewer is better — quality over quantity)
- Return [] if nothing worth remembering
- Do NOT extract: transient info (weather, greetings), obvious facts,
  information that was merely looked up (not personal to the user)
- DO extract: personal preferences, organizational context, deadlines,
  corrections to previous assumptions, working patterns
- Keys should be lowercase_snake_case, descriptive but short
- importance: 1-3 for minor preferences, 4-6 for useful context,
  7-9 for critical facts, 10 for safety-critical corrections
- Set supersedes when a fact updates or corrects an existing memory
"""


def _build_extraction_message(
    user_message: str,
    assistant_response: str,
    existing_keys: List[str],
) -> str:
    """Build the extraction prompt with conversation context."""
    parts = [
        f"User said: {user_message[:2000]}",
        f"\nAssistant responded: {assistant_response[:2000]}",
    ]

    if existing_keys:
        parts.append(
            f"\nExisting memory keys (set supersedes if updating one): "
            f"{', '.join(existing_keys[:50])}"
        )

    return "\n".join(parts)


# ════════════════════════════════════════════════════════════════
# Conversation Filter (skip trivial conversations)
# ════════════════════════════════════════════════════════════════

# Patterns that indicate a trivial conversation not worth extracting from
_TRIVIAL_PATTERNS = [
    re.compile(p, re.IGNORECASE)
    for p in [
        r"^(hi|hello|hey|thanks|thank you|ok|bye|goodbye)\s*[!.?]*$",
        r"^(yes|no|sure|cancel|stop|quit|nevermind)\s*[!.?]*$",
        r"^what (time|day) is it",
        r"^(help|menu|commands|skills)$",
    ]
]

# Minimum lengths to consider extraction
MIN_USER_MSG_LENGTH = 15
MIN_ASSISTANT_MSG_LENGTH = 50


def should_extract(user_message: str, assistant_response: str) -> bool:
    """
    Decide whether a conversation is worth extracting memories from.

    Trivial exchanges (greetings, yes/no, short lookups) are skipped
    to avoid wasting LLM calls on extraction.
    """
    # Too short
    if len(user_message) < MIN_USER_MSG_LENGTH:
        return False
    if len(assistant_response) < MIN_ASSISTANT_MSG_LENGTH:
        return False

    # Trivial patterns
    if any(p.search(user_message.strip()) for p in _TRIVIAL_PATTERNS):
        return False

    return True


# ════════════════════════════════════════════════════════════════
# Memory Extractor
# ════════════════════════════════════════════════════════════════


class MemoryExtractor:
    """
    Proactively extracts facts from conversations.

    Runs asynchronously after the agent response is sent, so
    the user never waits for extraction to complete.
    """

    def __init__(
        self,
        provider: Optional["LLMProvider"] = None,
        memory: Optional["Memory"] = None,
        enabled: bool = True,
        max_extraction_tokens: int = 1024,
    ):
        self.provider = provider
        self.memory = memory
        self.enabled = enabled
        self.max_extraction_tokens = max_extraction_tokens

        # Stats
        self.extractions_run = 0
        self.extractions_skipped = 0
        self.memories_stored = 0
        self.memories_superseded = 0
        self.total_extraction_time_ms = 0.0
        self._lock = threading.Lock()

    def extract_async(
        self,
        user_message: str,
        assistant_response: str,
        channel: str = "default",
    ):
        """
        Extract memories in a background thread. Fire-and-forget.

        Called after the response is sent to the user.
        """
        if not self.enabled or not self.provider or not self.memory:
            return

        if not should_extract(user_message, assistant_response):
            with self._lock:
                self.extractions_skipped += 1
            return

        thread = threading.Thread(
            target=self._extract_sync,
            args=(user_message, assistant_response, channel),
            daemon=True,
            name="memory-extractor",
        )
        thread.start()

    def _extract_sync(
        self,
        user_message: str,
        assistant_response: str,
        channel: str,
    ):
        """Synchronous extraction — runs in background thread."""
        start = time.time()

        try:
            existing_keys = list(self.memory.memories.keys())

            extraction_msg = _build_extraction_message(
                user_message,
                assistant_response,
                existing_keys,
            )

            response = self.provider.chat(
                messages=[{"role": "user", "content": extraction_msg}],
                tools=[],
                system=EXTRACTION_SYSTEM,
                max_tokens=self.max_extraction_tokens,
            )

            extracted = self._parse_extractions(response.text, existing_keys)

            # Store extracted memories
            for mem in extracted:
                # Handle supersedes: forget the old memory
                if mem.supersedes and mem.supersedes in self.memory.memories:
                    old_entry = self.memory.memories[mem.supersedes]
                    logger.info(
                        f"Memory superseded: {mem.supersedes} "
                        f"('{old_entry.value[:40]}') → '{mem.value[:40]}'"
                    )
                    self.memory.forget(mem.supersedes)
                    with self._lock:
                        self.memories_superseded += 1

                self.memory.remember(
                    key=mem.key,
                    value=mem.value,
                    category=mem.category,
                    source=f"extracted/{channel}",
                    importance=mem.importance,
                )
                with self._lock:
                    self.memories_stored += 1

            elapsed_ms = (time.time() - start) * 1000
            with self._lock:
                self.extractions_run += 1
                self.total_extraction_time_ms += elapsed_ms

            if extracted:
                logger.info(
                    f"Extracted {len(extracted)} memories ({elapsed_ms:.0f}ms): "
                    f"{', '.join(m.key for m in extracted)}"
                )
            else:
                logger.debug(f"No memories extracted ({elapsed_ms:.0f}ms)")

        except Exception as e:
            logger.warning(f"Memory extraction failed: {e}")

    def extract_sync(
        self,
        user_message: str,
        assistant_response: str,
        channel: str = "default",
    ) -> List[ExtractedMemory]:
        """
        Synchronous extraction for testing. Returns extracted memories.
        """
        if not self.provider or not self.memory:
            return []

        existing_keys = list(self.memory.memories.keys())

        extraction_msg = _build_extraction_message(
            user_message,
            assistant_response,
            existing_keys,
        )

        response = self.provider.chat(
            messages=[{"role": "user", "content": extraction_msg}],
            tools=[],
            system=EXTRACTION_SYSTEM,
            max_tokens=self.max_extraction_tokens,
        )

        return self._parse_extractions(response.text, existing_keys)

    def _parse_extractions(
        self,
        response_text: str,
        existing_keys: List[str],
    ) -> List[ExtractedMemory]:
        """Parse the LLM's JSON array response into ExtractedMemory objects."""
        if not response_text:
            return []

        # Strip markdown fences
        text = response_text.strip()
        text = re.sub(r"^```(?:json)?\s*", "", text)
        text = re.sub(r"\s*```$", "", text)

        # Parse JSON
        try:
            data = json.loads(text)
        except json.JSONDecodeError:
            # Try to find JSON array
            match = re.search(r"\[[\s\S]*\]", text)
            if match:
                try:
                    data = json.loads(match.group())
                except json.JSONDecodeError:
                    return []
            else:
                return []

        if not isinstance(data, list):
            return []

        # Validate and convert
        results = []
        existing_set = set(existing_keys)

        for item in data[:5]:  # Cap at 5
            if not isinstance(item, dict):
                continue

            key = item.get("key", "").strip()
            value = item.get("value", "").strip()
            category = item.get("category", "fact").strip()

            if not key or not value:
                continue

            # Sanitize key
            key = re.sub(r"[^a-z0-9_]", "_", key.lower())[:50]

            # Validate category
            valid_cats = {"user_info", "org_info", "deadline", "preference", "correction", "fact"}
            if category not in valid_cats:
                category = "fact"

            # Validate importance
            importance = item.get("importance", 5)
            if not isinstance(importance, int):
                importance = 5
            importance = max(1, min(10, importance))

            # Validate supersedes
            supersedes = item.get("supersedes")
            if supersedes and supersedes not in existing_set:
                supersedes = None  # Can't supersede a non-existent key

            # Skip if exact duplicate of existing memory
            if key in existing_set:
                existing = self.memory.memories.get(key)
                if existing and existing.value == value:
                    continue  # Exact duplicate, skip

            results.append(
                ExtractedMemory(
                    key=key,
                    value=value,
                    category=category,
                    importance=importance,
                    supersedes=supersedes,
                )
            )

        return results

    def get_stats(self) -> Dict[str, Any]:
        """Return extraction statistics."""
        with self._lock:
            total = self.extractions_run + self.extractions_skipped
            return {
                "extractions_run": self.extractions_run,
                "extractions_skipped": self.extractions_skipped,
                "total_conversations": total,
                "memories_stored": self.memories_stored,
                "memories_superseded": self.memories_superseded,
                "extraction_rate": (self.extractions_run / total if total > 0 else 0),
                "avg_extraction_time_ms": (
                    self.total_extraction_time_ms / self.extractions_run
                    if self.extractions_run > 0
                    else 0
                ),
            }


# ════════════════════════════════════════════════════════════════
# Memory Decay
# ════════════════════════════════════════════════════════════════


def decay_memories(
    memory: "Memory",
    days_idle: int = 30,
    decay_amount: int = 1,
    archive_threshold: int = 2,
    archive_after_days: int = 90,
) -> Dict[str, Any]:
    """
    Reduce importance of idle memories and archive very old ones.

    Call periodically (e.g., daily) to keep memory fresh.

    Args:
        memory: Memory instance to decay
        days_idle: Days since last update before decay starts
        decay_amount: How much to reduce importance per cycle
        archive_threshold: Archive if importance drops to or below this
        archive_after_days: Only archive if also older than this many days

    Returns:
        Stats dict with counts of decayed and archived memories
    """
    now = datetime.now()
    decayed = 0
    archived = []

    for key, entry in list(memory.memories.items()):
        try:
            updated = datetime.fromisoformat(entry.updated_at)
        except (ValueError, TypeError) as e:
            logger.debug("Value conversion error suppressed: %s", e)
            continue
        days_since_update = (now - updated).days

        if days_since_update < days_idle:
            continue

        # Decay importance
        new_importance = max(1, entry.importance - decay_amount)
        if new_importance != entry.importance:
            entry.importance = new_importance
            decayed += 1

        # Archive check: low importance AND very old
        if entry.importance <= archive_threshold and days_since_update >= archive_after_days:
            archived.append(key)

    # Archive (remove from active memory)
    for key in archived:
        logger.info(f"Archiving memory: {key} (importance={memory.memories[key].importance})")
        memory.forget(key)

    # Save if we decayed anything
    if decayed > 0:
        memory._save()
        if memory._vectors_dirty:
            memory._save_vectors()

    stats = {
        "decayed": decayed,
        "archived": len(archived),
        "archived_keys": archived,
        "total_memories": len(memory.memories),
    }

    if decayed > 0 or archived:
        logger.info(
            f"Memory decay: {decayed} decayed, {len(archived)} archived, "
            f"{len(memory.memories)} remaining"
        )

    return stats
