import json
import os
import hashlib
from datetime import datetime
from typing import List, Dict, Any, Optional
from pathlib import Path

class JSONLibrary:
    def __init__(self, base_path: str = "./data/json_library/"):
        self.base_path = Path(base_path)
        self.base_path.mkdir(parents=True, exist_ok=True)
        
        self.categories = [
            "user_preferences",
            "project_context",
            "conversation_history",
            "important_facts",
            "code_patterns",
            "learned_skills"
        ]
        
        self._init_categories()
    
    def _init_categories(self):
        for category in self.categories:
            category_path = self.base_path / f"{category}.json"
            if not category_path.exists():
                with open(category_path, "w") as f:
                    json.dump([], f)
    
    def _get_category_path(self, category: str) -> Path:
        if category not in self.categories:
            category = "important_facts"
        return self.base_path / f"{category}.json"
    
    MAX_ENTRIES_PER_CATEGORY = 1000

    def store(self, content: Dict[str, Any], category: str = "important_facts") -> str:
        path = self._get_category_path(category)

        with open(path, "r") as f:
            data = json.load(f)

        # Trim oldest entries to stay within cap
        if len(data) >= self.MAX_ENTRIES_PER_CATEGORY:
            data = data[-(self.MAX_ENTRIES_PER_CATEGORY - 1):]

        memory_id = hashlib.md5(f"{datetime.now().isoformat()}{content}".encode()).hexdigest()[:8]

        entry = {
            "id": memory_id,
            "timestamp": datetime.now().isoformat(),
            "content": content,
            "category": category,
            "access_count": 0,
            "last_accessed": None
        }

        data.append(entry)

        with open(path, "w") as f:
            json.dump(data, f, indent=2)

        return memory_id
    
    def recall(self, query: Optional[str] = None, category: Optional[str] = None, limit: int = 10) -> List[Dict]:
        cats = [category] if (category and category in self.categories) else self.categories

        # Load all relevant files, keeping track of origin for write-back
        file_data: Dict[str, list] = {}
        all_memories = []
        for cat in cats:
            path = self._get_category_path(cat)
            if path.exists():
                with open(path, "r") as f:
                    data = json.load(f)
                file_data[cat] = data
                all_memories.extend(data)

        all_memories.sort(key=lambda x: x["timestamp"], reverse=True)

        if query:
            query_lower = query.lower()
            all_memories = [m for m in all_memories if query_lower in json.dumps(m["content"]).lower()]

        result = all_memories[:limit]

        # Update access metadata and write back to disk
        accessed_ids = {m["id"] for m in result}
        now = datetime.now().isoformat()
        for cat, data in file_data.items():
            modified = False
            for entry in data:
                if entry["id"] in accessed_ids:
                    entry["access_count"] += 1
                    entry["last_accessed"] = now
                    modified = True
            if modified:
                with open(self._get_category_path(cat), "w") as f:
                    json.dump(data, f, indent=2)

        return result
    
    def get_context_string(self, limit: int = 5) -> str:
        recent = self.recall(limit=limit)
        if not recent:
            return ""
        
        context = "\n## SHOREKEEPER Memory\n"
        context += "The following information should be remembered:\n\n"
        
        for mem in recent:
            content = mem["content"]
            if isinstance(content, dict):
                for key, value in content.items():
                    context += f"- {key}: {value}\n"
            else:
                context += f"- {content}\n"
        
        return context