"""
KIESSCLAW Memory Engine
File-based persistent memory system for all agents.
Each agent has its own memory namespace. Shared workspace brain for cross-agent context.
"""

from __future__ import annotations
import json
import logging
from datetime import datetime
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)


class MemoryEngine:
    """
    Persistent markdown-based memory for KIESSCLAW agents.

    Each agent reads/writes to its own namespace.
    Shared memory available for cross-agent coordination.
    """

    def __init__(self, workspace_dir: str = "./workspace"):
        self.workspace = Path(workspace_dir)
        self.memory_dir = self.workspace / "memory"
        self.data_dir = self.workspace / "data"
        self.memory_dir.mkdir(parents=True, exist_ok=True)
        self.data_dir.mkdir(parents=True, exist_ok=True)
        self._required_collections = {
            "accounts",
            "contacts",
            "segments",
            "sequences",
            "enrollments",
            "messages",
            "replies",
            "analytics",
            "seo_audits",
            "waitlist",
            "scheduled_jobs",
        }
        self._ensure_default_collections()

    # -------------------------------------------------------------------------
    # Core read/write
    # -------------------------------------------------------------------------

    def read(self, agent: str, key: str) -> Optional[str]:
        """Read a memory file for a given agent and key."""
        path = self._path(agent, key)
        if path.exists():
            return path.read_text(encoding="utf-8")
        return None

    def write(self, agent: str, key: str, content: str) -> None:
        """Write/overwrite a memory file."""
        path = self._path(agent, key)
        path.write_text(content, encoding="utf-8")
        logger.debug(f"[Memory] Wrote {agent}/{key}")

    def append(self, agent: str, key: str, content: str) -> None:
        """Append to a memory file (useful for logs)."""
        path = self._path(agent, key)
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        with path.open("a", encoding="utf-8") as f:
            f.write(f"\n\n---\n**{timestamp}**\n{content}")
        logger.debug(f"[Memory] Appended to {agent}/{key}")

    def delete(self, agent: str, key: str) -> bool:
        """Delete a memory file."""
        path = self._path(agent, key)
        if path.exists():
            path.unlink()
            return True
        return False

    def list_keys(self, agent: str) -> list[str]:
        """List all memory keys for an agent."""
        agent_dir = self.memory_dir / agent
        if not agent_dir.exists():
            return []
        return [f.stem for f in agent_dir.glob("*.md")]

    # -------------------------------------------------------------------------
    # Shared / Cross-agent memory
    # -------------------------------------------------------------------------

    def read_shared(self, key: str) -> Optional[str]:
        return self.read("shared", key)

    def write_shared(self, key: str, content: str) -> None:
        self.write("shared", key, content)

    def append_shared(self, key: str, content: str) -> None:
        self.append("shared", key, content)

    # -------------------------------------------------------------------------
    # Log helpers
    # -------------------------------------------------------------------------

    def log(self, agent: str, message: str) -> None:
        """Append a timestamped log entry for an agent."""
        self.append(agent, "log", message)

    def log_prospect(self, domain: str, signal: str, score: int) -> None:
        """Log a prospect signal to shared memory."""
        entry = f"**Domain**: {domain}  \n**Signal**: {signal}  \n**ICP Score**: {score}/100"
        self.append("shared", "prospects", entry)
        logger.info(f"[Prospect Signal] {domain} - Score: {score}")

    def log_reply(self, contact_email: str, intent: str, confidence: float) -> None:
        """Log a reply classification to shared memory."""
        entry = f"**Contact**: {contact_email}  \n**Intent**: {intent}  \n**Confidence**: {confidence:.2f}"
        self.append("shared", "replies", entry)
        logger.info(f"[Reply] {contact_email} - Intent: {intent}")

    # -------------------------------------------------------------------------
    # State management (JSON-serialized)
    # -------------------------------------------------------------------------

    def get_state(self, agent: str) -> dict:
        """Get JSON state for an agent."""
        path = self._state_path(agent)
        if path.exists():
            try:
                return json.loads(path.read_text(encoding="utf-8"))
            except json.JSONDecodeError:
                return {}
        return {}

    def set_state(self, agent: str, state: dict) -> None:
        """Save JSON state for an agent."""
        path = self._state_path(agent)
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(state, indent=2), encoding="utf-8")

    def update_state(self, agent: str, updates: dict) -> None:
        """Merge updates into existing agent state."""
        state = self.get_state(agent)
        state.update(updates)
        self.set_state(agent, state)

    # -------------------------------------------------------------------------
    # Data storage (JSON files in data directory)
    # -------------------------------------------------------------------------

    def load_data(self, collection: str) -> list[dict]:
        """Load a JSON data collection (e.g., 'contacts', 'sequences')."""
        path = self.data_dir / f"{collection}.json"
        if not path.exists():
            self.save_data(collection, [])
            return []

        try:
            data = json.loads(path.read_text(encoding="utf-8"))
        except json.JSONDecodeError as exc:
            raise RuntimeError(f"Invalid JSON in data collection: {path}") from exc

        if not isinstance(data, list):
            raise RuntimeError(f"Collection must be a JSON list: {path}")

        return data

    def save_data(self, collection: str, data: list[dict]) -> None:
        """Save a JSON data collection."""
        path = self.data_dir / f"{collection}.json"
        path.write_text(json.dumps(data, indent=2, default=str), encoding="utf-8")
        logger.debug(f"[Data] Saved {len(data)} items to {collection}")

    def append_data(self, collection: str, item: dict) -> None:
        """Append a single item to a data collection."""
        data = self.load_data(collection)
        data.append(item)
        self.save_data(collection, data)

    # -------------------------------------------------------------------------
    # Private helpers
    # -------------------------------------------------------------------------

    def _path(self, agent: str, key: str) -> Path:
        agent_dir = self.memory_dir / agent
        agent_dir.mkdir(parents=True, exist_ok=True)
        return agent_dir / f"{key}.md"

    def _state_path(self, agent: str) -> Path:
        return self.memory_dir / agent / "_state.json"

    def _ensure_default_collections(self) -> None:
        """Ensure core workspace collections always exist on disk."""
        for collection in self._required_collections:
            path = self.data_dir / f"{collection}.json"
            if not path.exists():
                self.save_data(collection, [])
