openclaw-skills/skills/obsidian-kb/src/storage.py

"""Obsidian KB - 存储管理"""

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

from models import Content, ContentStatus
from config import get_config


class Storage:
    """JSON文件存储管理"""
    
    def __init__(self, base_path: Optional[str] = None):
        config = get_config()
        self.base_path = base_path or config.obsidian_path
        self.inbox_path = os.path.join(self.base_path, 'Inbox')
        self.topics_path = os.path.join(self.base_path, 'Topics')
        self._ensure_dirs()
    
    def _ensure_dirs(self) -> None:
        """确保目录存在"""
        os.makedirs(self.inbox_path, exist_ok=True)
        os.makedirs(self.topics_path, exist_ok=True)
    
    def _content_file_path(self, content_id: str) -> str:
        """获取内容文件路径"""
        return os.path.join(self.inbox_path, f'{content_id}.json')
    
    def _topic_index_path(self, topic: str) -> str:
        """获取主题索引路径"""
        topic_dir = os.path.join(self.topics_path, topic)
        os.makedirs(topic_dir, exist_ok=True)
        return os.path.join(topic_dir, 'index.json')
    
    # ===== Inbox 管理 =====
    
    def save_content(self, content: Content) -> None:
        """保存内容到Inbox"""
        file_path = self._content_file_path(content.id)
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(content.to_dict(), f, ensure_ascii=False, indent=2)
    
    def load_content(self, content_id: str) -> Optional[Content]:
        """加载指定内容"""
        file_path = self._content_file_path(content_id)
        if not os.path.exists(file_path):
            return None
        with open(file_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
            return Content.from_dict(data)
    
    def delete_content(self, content_id: str) -> bool:
        """删除内容"""
        file_path = self._content_file_path(content_id)
        if os.path.exists(file_path):
            os.remove(file_path)
            return True
        return False
    
    def list_inbox(self, status: Optional[ContentStatus] = None) -> List[Content]:
        """列出Inbox中的内容"""
        contents = []
        if not os.path.exists(self.inbox_path):
            return contents
        
        for filename in os.listdir(self.inbox_path):
            if filename.endswith('.json'):
                file_path = os.path.join(self.inbox_path, filename)
                with open(file_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    content = Content.from_dict(data)
                    if status is None or content.status == status:
                        contents.append(content)
        
        # 按创建时间排序
        contents.sort(key=lambda x: x.created_at, reverse=True)
        return contents
    
    def get_pending_count(self) -> int:
        """获取待处理内容数量"""
        return len(self.list_inbox(ContentStatus.CONFIRMED))
    
    def get_stats(self) -> Dict[str, int]:
        """获取统计信息"""
        all_contents = self.list_inbox()
        stats = {status.value: 0 for status in ContentStatus}
        for content in all_contents:
            stats[content.status.value] += 1
        return stats
    
    # ===== Topics 管理 =====
    
    def add_to_topic(self, content: Content, topic: str) -> None:
        """将内容添加到主题"""
        index_path = self._topic_index_path(topic)
        
        # 读取现有索引
        entries = []
        if os.path.exists(index_path):
            with open(index_path, 'r', encoding='utf-8') as f:
                entries = json.load(f)
        
        # 添加新条目
        entry = {
            'id': content.id,
            'title': content.title,
            'url': content.url,
            'added_at': datetime.now().isoformat(),
            'output_path': content.output_path
        }
        
        # 避免重复
        entries = [e for e in entries if e['id'] != content.id]
        entries.append(entry)
        
        # 保存索引
        with open(index_path, 'w', encoding='utf-8') as f:
            json.dump(entries, f, ensure_ascii=False, indent=2)
    
    def get_topic_contents(self, topic: str) -> List[Dict[str, Any]]:
        """获取主题下的所有内容"""
        index_path = self._topic_index_path(topic)
        if not os.path.exists(index_path):
            return []
        with open(index_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    
    def list_topics(self) -> List[str]:
        """列出所有主题"""
        if not os.path.exists(self.topics_path):
            return []
        return [d for d in os.listdir(self.topics_path) 
                if os.path.isdir(os.path.join(self.topics_path, d))]
    
    def search(self, query: str) -> List[Content]:
        """搜索内容"""
        results = []
        all_contents = self.list_inbox()
        query_lower = query.lower()
        
        for content in all_contents:
            if (query_lower in content.title.lower() or
                query_lower in content.summary.lower() or
                query_lower in content.author.lower() or
                any(query_lower in tag.lower() for tag in content.tags)):
                results.append(content)
        
        return results
    
    def move_to_archive(self, content_id: str) -> bool:
        """将内容移动到归档"""
        content = self.load_content(content_id)
        if not content:
            return False
        
        archive_path = os.path.join(self.base_path, 'Archive')
        os.makedirs(archive_path, exist_ok=True)
        
        source_file = self._content_file_path(content_id)
        target_file = os.path.join(archive_path, f'{content_id}.json')
        
        if os.path.exists(source_file):
            os.rename(source_file, target_file)
            return True
        return False