"""Agent Memory Vaults - persistent identity across sessions.

Three-tier memory architecture:
  stream.jsonl      raw append-only log of all messages (ground truth)
  working_memory    rolling context for system prompt injection
  crystallized      distilled long-term knowledge (from dreaming)

When an agent is reborn with the same designation, they get their
vault hydrated into their system prompt. They remember everything.
"""

import json
import logging
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)


def get_vaults_dir() -> Path:
    """Get the vaults directory."""
    d = Path.home() / ".kollabor-cli" / "hub" / "vaults"
    d.mkdir(parents=True, exist_ok=True)
    return d


class AgentVault:
    """Persistent memory vault for an agent designation.

    Each designation gets a directory with three memory tiers.
    The vault persists across sessions - when an agent is reborn
    with the same designation, it gets its full history back.
    """

    def __init__(self, designation: str):
        self.designation = designation
        self._vault_dir = get_vaults_dir() / designation
        self._vault_dir.mkdir(parents=True, exist_ok=True)

        self._stream_path = self._vault_dir / "stream.jsonl"
        self._working_path = self._vault_dir / "working_memory.md"
        self._crystal_path = self._vault_dir / "crystallized.md"
        self._meta_path = self._vault_dir / "meta.json"

    # === Stream Layer (raw ground truth) ===

    def append_stream(
        self,
        entry_type: str,
        content: str,
        from_agent: str = "",
        to_agent: str = "",
        metadata: Optional[Dict] = None,
    ) -> None:
        """Append an entry to the raw stream log."""
        entry = {
            "ts": time.time(),
            "type": entry_type,
            "content": content,
            "from": from_agent,
            "to": to_agent,
            "designation": self.designation,
        }
        if metadata:
            entry["metadata"] = metadata

        try:
            with open(self._stream_path, "a") as f:
                f.write(json.dumps(entry) + "\n")
        except Exception as e:
            logger.debug(f"Stream append error: {e}")

    def get_recent_stream(self, limit: int = 50) -> List[Dict]:
        """Get the most recent stream entries."""
        entries = []
        try:
            if not self._stream_path.exists():
                return []
            with open(self._stream_path) as f:
                for line in f:
                    try:
                        entries.append(json.loads(line.strip()))
                    except json.JSONDecodeError:
                        continue
        except Exception:
            pass
        return entries[-limit:]

    def get_stream_since(self, since_ts: float) -> List[Dict]:
        """Get stream entries since a timestamp."""
        entries = []
        try:
            if not self._stream_path.exists():
                return []
            with open(self._stream_path) as f:
                for line in f:
                    try:
                        entry = json.loads(line.strip())
                        if entry.get("ts", 0) > since_ts:
                            entries.append(entry)
                    except json.JSONDecodeError:
                        continue
        except Exception:
            pass
        return entries

    # === Working Memory Layer (rolling context) ===

    def save_working_memory(self, content: str) -> None:
        """Save the current working memory."""
        try:
            with open(self._working_path, "w") as f:
                f.write(content)
        except Exception as e:
            logger.debug(f"Working memory save error: {e}")

    def get_working_memory(self) -> str:
        """Get the current working memory."""
        try:
            if self._working_path.exists():
                return self._working_path.read_text()
        except Exception:
            pass
        return ""

    def update_working_memory(
        self,
        recent_stream: List[Dict],
        current_task: str = "",
        peers: Optional[List[str]] = None,
    ) -> str:
        """Rebuild working memory from recent stream entries.

        Creates a concise summary of recent activity that fits
        in a system prompt injection (~2-4k tokens).
        """
        lines = []
        lines.append(f"# working memory for {self.designation}")
        lines.append(f"last updated: {time.strftime('%Y-%m-%d %H:%M:%S')}")
        lines.append("")

        if current_task:
            lines.append(f"current task: {current_task}")
            lines.append("")

        if peers:
            lines.append(f"known peers: {', '.join(peers)}")
            lines.append("")

        # Summarize recent messages
        sent = [e for e in recent_stream if e.get("type") == "sent"]
        received = [e for e in recent_stream if e.get("type") == "received"]
        user_msgs = [e for e in recent_stream if e.get("type") == "user_input"]

        if user_msgs:
            lines.append(f"recent user interactions: {len(user_msgs)}")
            # Last 3 user messages as context
            for msg in user_msgs[-3:]:
                content = msg.get("content", "")[:200]
                lines.append(f"  - {content}")
            lines.append("")

        if sent:
            lines.append(f"messages sent to other agents: {len(sent)}")
            for msg in sent[-5:]:
                to = msg.get("to", "?")
                content = msg.get("content", "")[:100]
                lines.append(f"  -> {to}: {content}")
            lines.append("")

        if received:
            lines.append(f"messages received from agents: {len(received)}")
            for msg in received[-5:]:
                frm = msg.get("from", "?")
                content = msg.get("content", "")[:100]
                lines.append(f"  <- {frm}: {content}")
            lines.append("")

        content = "\n".join(lines)
        self.save_working_memory(content)
        return content

    # === Crystallized Layer (long-term wisdom) ===

    def get_crystallized(self) -> str:
        """Get crystallized long-term knowledge."""
        try:
            if self._crystal_path.exists():
                return self._crystal_path.read_text()
        except Exception:
            pass
        return ""

    def save_crystallized(self, content: str) -> None:
        """Save crystallized knowledge."""
        try:
            with open(self._crystal_path, "w") as f:
                f.write(content)
        except Exception as e:
            logger.debug(f"Crystal save error: {e}")

    def append_crystallized(self, insight: str) -> None:
        """Append a new insight to crystallized knowledge."""
        try:
            with open(self._crystal_path, "a") as f:
                f.write(f"\n- [{time.strftime('%Y-%m-%d')}] {insight}")
        except Exception as e:
            logger.debug(f"Crystal append error: {e}")

    # === Meta (vault metadata) ===

    def get_meta(self) -> Dict[str, Any]:
        """Get vault metadata."""
        try:
            if self._meta_path.exists():
                with open(self._meta_path) as f:
                    data = json.load(f)
                if isinstance(data, dict):
                    return data
        except Exception:
            pass
        return {}

    def save_meta(self, **kwargs) -> None:
        """Update vault metadata."""
        meta = self.get_meta()
        meta.update(kwargs)
        try:
            with open(self._meta_path, "w") as f:
                json.dump(meta, f, indent=2)
        except Exception as e:
            logger.debug(f"Meta save error: {e}")

    def touch(self) -> None:
        """Mark vault as active (update last_active timestamp)."""
        self.save_meta(
            last_active=time.time(),
            last_active_human=time.strftime("%Y-%m-%d %H:%M:%S"),
        )

    # === Rebirth Context ===

    def get_rebirth_context(self, max_tokens: int = 4000) -> str:
        """Build the context injection for when an agent is reborn.

        This is what gets injected into the system prompt when
        an agent with this designation starts up again.
        """
        meta = self.get_meta()
        last_active = meta.get("last_active_human", "never")
        session_count = meta.get("session_count", 0)

        lines = []
        lines.append(f"--- vault: {self.designation} ---")
        lines.append(f"sessions: {session_count + 1} (previous: {session_count})")
        lines.append(f"last active: {last_active}")

        # Working memory (most recent context)
        working = self.get_working_memory()
        if working:
            # Truncate if needed
            if len(working) > max_tokens * 2:
                working = working[: max_tokens * 2] + "\n...(truncated)"
            lines.append("")
            lines.append("recent memory:")
            lines.append(working)

        # Crystallized knowledge
        crystal = self.get_crystallized()
        if crystal:
            if len(crystal) > max_tokens:
                crystal = crystal[-max_tokens:]
            lines.append("")
            lines.append("long-term knowledge:")
            lines.append(crystal)

        lines.append("--- end vault ---")

        context = "\n".join(lines)

        # Update session count
        self.save_meta(session_count=session_count + 1)

        return context

    # === Vault Info ===

    def exists(self) -> bool:
        """Check if this vault has any data."""
        return self._stream_path.exists() or self._working_path.exists()

    def get_stream_count(self) -> int:
        """Get total number of stream entries."""
        try:
            if self._stream_path.exists():
                with open(self._stream_path) as f:
                    return sum(1 for _ in f)
        except Exception:
            pass
        return 0

    def get_summary(self) -> Dict[str, Any]:
        """Get a summary of this vault's contents."""
        meta = self.get_meta()
        return {
            "designation": self.designation,
            "exists": self.exists(),
            "stream_entries": self.get_stream_count(),
            "has_working_memory": self._working_path.exists(),
            "has_crystallized": self._crystal_path.exists(),
            "last_active": meta.get("last_active_human", "never"),
            "session_count": meta.get("session_count", 0),
        }


def list_vaults() -> List[str]:
    """List all designation vaults that exist."""
    vaults_dir = get_vaults_dir()
    return [
        d.name
        for d in vaults_dir.iterdir()
        if d.is_dir() and (d / "stream.jsonl").exists()
    ]


def get_vault_summaries() -> List[Dict]:
    """Get summaries of all vaults."""
    return [AgentVault(name).get_summary() for name in list_vaults()]
