"""
经验记忆模块 - 存储爬虫运行经验和教训
"""

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


class Memory:
    """经验记忆管理器"""
    
    def __init__(self, memory_dir: str = "./memory"):
        self.memory_dir = Path(memory_dir)
        self.memory_dir.mkdir(parents=True, exist_ok=True)
        self.memory_file = self.memory_dir / "experiences.json"
        self._load()
    
    def _load(self):
        """加载记忆"""
        if self.memory_file.exists():
            with open(self.memory_file, "r", encoding="utf-8") as f:
                self.data = json.load(f)
        else:
            self.data = {
                "websites": {},  # 按网站分类的经验
                "tasks": {},     # 按任务类型分类的经验
                "general": [],   # 通用经验
            }
    
    def _save(self):
        """保存记忆"""
        with open(self.memory_file, "w", encoding="utf-8") as f:
            json.dump(self.data, f, ensure_ascii=False, indent=2)
    
    def add_website_experience(
        self, 
        website: str, 
        experience: str,
        category: str = "general"
    ):
        """添加网站相关经验"""
        if website not in self.data["websites"]:
            self.data["websites"][website] = {}
        
        if category not in self.data["websites"][website]:
            self.data["websites"][website][category] = []
        
        entry = {
            "content": experience,
            "timestamp": datetime.now().isoformat(),
        }
        self.data["websites"][website][category].append(entry)
        self._save()
    
    def add_task_experience(
        self, 
        task_type: str, 
        experience: str
    ):
        """添加任务相关经验"""
        if task_type not in self.data["tasks"]:
            self.data["tasks"][task_type] = []
        
        entry = {
            "content": experience,
            "timestamp": datetime.now().isoformat(),
        }
        self.data["tasks"][task_type].append(entry)
        self._save()
    
    def add_general_experience(self, experience: str):
        """添加通用经验"""
        entry = {
            "content": experience,
            "timestamp": datetime.now().isoformat(),
        }
        self.data["general"].append(entry)
        self._save()
    
    def get_website_experiences(
        self, 
        website: str, 
        category: Optional[str] = None
    ) -> List[str]:
        """获取网站相关经验"""
        if website not in self.data["websites"]:
            return []
        
        if category:
            return [
                e["content"] 
                for e in self.data["websites"][website].get(category, [])
            ]
        
        # 返回所有分类的经验
        experiences = []
        for cat_experiences in self.data["websites"][website].values():
            experiences.extend([e["content"] for e in cat_experiences])
        return experiences
    
    def get_task_experiences(self, task_type: str) -> List[str]:
        """获取任务相关经验"""
        return [
            e["content"] 
            for e in self.data["tasks"].get(task_type, [])
        ]
    
    def get_general_experiences(self) -> List[str]:
        """获取通用经验"""
        return [e["content"] for e in self.data["general"]]
    
    def get_all_for_context(
        self, 
        website: Optional[str] = None,
        task_type: Optional[str] = None
    ) -> str:
        """获取所有相关经验，用于构建上下文"""
        lines = []
        
        # 通用经验
        general = self.get_general_experiences()
        if general:
            lines.append("【通用经验】")
            for exp in general:
                lines.append(f"  - {exp}")
        
        # 网站经验
        if website:
            site_exp = self.get_website_experiences(website)
            if site_exp:
                lines.append(f"\n【{website} 网站经验】")
                for exp in site_exp:
                    lines.append(f"  - {exp}")
        
        # 任务经验
        if task_type:
            task_exp = self.get_task_experiences(task_type)
            if task_exp:
                lines.append(f"\n【{task_type} 任务经验】")
                for exp in task_exp:
                    lines.append(f"  - {exp}")
        
        return "\n".join(lines) if lines else ""
    
    def print_all(self):
        """打印所有经验"""
        print("\n" + "=" * 50)
        print("  经验记忆库")
        print("=" * 50)
        
        # 通用经验
        if self.data["general"]:
            print("\n【通用经验】")
            for e in self.data["general"]:
                print(f"  - {e['content']}")
        
        # 网站经验
        if self.data["websites"]:
            for site, categories in self.data["websites"].items():
                print(f"\n【{site}】")
                for cat, exps in categories.items():
                    print(f"  [{cat}]")
                    for e in exps:
                        print(f"    - {e['content']}")
        
        # 任务经验
        if self.data["tasks"]:
            for task, exps in self.data["tasks"].items():
                print(f"\n【任务: {task}】")
                for e in exps:
                    print(f"  - {e['content']}")


# 全局记忆实例
_memory: Optional[Memory] = None


def get_memory(memory_dir: str = "./memory") -> Memory:
    """获取全局记忆实例"""
    global _memory
    if _memory is None:
        _memory = Memory(memory_dir)
    return _memory
