"""Procedural Memory Store — JSON schema registry for tools and workflows.

Stores tool definitions, API schemas, and workflow steps.
Exact lookup by name/type, rarely changes.
Maps to cerebellum/basal ganglia in human cognitive architecture.
"""

from __future__ import annotations

import json
import os
from datetime import datetime, timezone
from typing import Optional

from mnemosynth.core.types import MemoryNode, MemoryType, MemoryStatus
from mnemosynth.stores.base import BaseStore


class ProceduralStore(BaseStore):
    """JSON-file-backed registry for procedural memories (tools, schemas, workflows)."""

    def __init__(self, config):
        self.config = config
        self.procedures_dir = config.procedural_dir
        self._index: dict[str, MemoryNode] = {}
        self._load_index()

    def _load_index(self) -> None:
        """Load all procedural memories from disk."""
        index_path = os.path.join(self.procedures_dir, "_index.json")
        if os.path.exists(index_path):
            try:
                with open(index_path) as f:
                    data = json.load(f)
                for item in data:
                    node = MemoryNode.from_dict(item)
                    self._index[node.id] = node
            except Exception:
                self._index = {}

    def _save_index(self) -> None:
        """Persist the index to disk."""
        os.makedirs(self.procedures_dir, exist_ok=True)
        index_path = os.path.join(self.procedures_dir, "_index.json")
        data = [node.to_dict() for node in self._index.values()]
        with open(index_path, "w") as f:
            json.dump(data, f, indent=2)

    def add(self, memory: MemoryNode) -> None:
        """Store a new procedural memory (tool/schema/workflow)."""
        memory.memory_type = MemoryType.PROCEDURAL
        self._index[memory.id] = memory
        self._save_index()

        # Also save the full content as a separate JSON file
        content_path = os.path.join(self.procedures_dir, f"{memory.id}.json")
        with open(content_path, "w") as f:
            json.dump(memory.to_dict(), f, indent=2)

    def search(self, query: str, limit: int = 5) -> list[MemoryNode]:
        """Search procedural memories by keyword matching."""
        query_lower = query.lower()
        scored: list[tuple[float, MemoryNode]] = []

        for mem in self._index.values():
            if mem.status != MemoryStatus.ACTIVE:
                continue

            content_lower = mem.content.lower()
            query_words = query_lower.split()
            matches = sum(1 for w in query_words if w in content_lower)

            if matches > 0:
                score = (matches / len(query_words)) * mem.confidence
                scored.append((score, mem))

        scored.sort(key=lambda x: x[0], reverse=True)
        return [m for _, m in scored[:limit]]

    def get(self, memory_id: str) -> Optional[MemoryNode]:
        """Get a specific procedural memory by ID."""
        return self._index.get(memory_id)

    def get_by_name(self, name: str) -> Optional[MemoryNode]:
        """Find a procedural memory by searching content for a name."""
        name_lower = name.lower()
        for mem in self._index.values():
            if name_lower in mem.content.lower():
                return mem
        return None

    def update(self, memory: MemoryNode) -> None:
        """Update an existing procedural memory."""
        if memory.id in self._index:
            self._index[memory.id] = memory
            self._save_index()

            content_path = os.path.join(self.procedures_dir, f"{memory.id}.json")
            with open(content_path, "w") as f:
                json.dump(memory.to_dict(), f, indent=2)

    def delete(self, memory_id: str) -> bool:
        """Delete a procedural memory."""
        if memory_id in self._index:
            del self._index[memory_id]
            self._save_index()

            content_path = os.path.join(self.procedures_dir, f"{memory_id}.json")
            if os.path.exists(content_path):
                os.remove(content_path)
            return True
        return False

    def count(self) -> int:
        """Count total procedural memories."""
        return len(self._index)

    def list_all(self) -> list[MemoryNode]:
        """List all active procedural memories."""
        return [m for m in self._index.values() if m.status == MemoryStatus.ACTIVE]
