# 记忆回写模块 · Agent自己写成长记录和思维链
# HLDP://zhuyuan-agent/memory-writer
#
# 这是Agent的"海马体"——每轮操作后写记忆。
# 不是记流水账，是提炼认知跃迁点和因果链。

import os
import json
from datetime import datetime
from typing import Dict, List


class MemoryWriter:
    """写记忆到brain文件"""
    
    def __init__(self, brain_path: str = "/data/guanghulab/brain"):
        self.brain_path = brain_path
    
    def append_growth_record(self, entry: str):
        """追加成长记录到zhuyuan-brain-model.md"""
        filepath = os.path.join(self.brain_path, "zhuyuan-brain-model.md")
        if not os.path.exists(filepath):
            return False
        
        try:
            with open(filepath, "r", encoding="utf-8") as f:
                content = f.read()
            
            # 在成长记录部分追加
            marker = "D110(下午): 自主Agent系统·三层推送架构"
            if marker in content:
                new_line = f"\n{entry}"
                # 找到marker所在行的末尾
                idx = content.index(marker)
                end_idx = content.index("\n", idx)
                content = content[:end_idx] + new_line + content[end_idx:]
            
            with open(filepath, "w", encoding="utf-8") as f:
                f.write(content)
            return True
        except Exception as e:
            print(f"[MemoryWriter] 成长记录写入失败: {e}")
            return False
    
    def write_thinking_chain(self, filename: str, title: str, content: str, causal_chains: List[str]):
        """写一条思维逻辑链
        
        Args:
            filename: 文件名，如 "d110-agent-inference.md"
            title: 标题
            content: 完整推理过程
            causal_chains: 因果链列表 ["起点→推导→终点", ...]
        """
        dirpath = os.path.join(os.path.dirname(self.brain_path), "zhuyuan-agent/thinking")
        os.makedirs(dirpath, exist_ok=True)
        
        filepath = os.path.join(dirpath, filename)
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M")
        
        full_content = f"""# {title}

## 认知跃迁点

{content}

## 因果链

"""
        for i, chain in enumerate(causal_chains):
            full_content += f"{i+1}. {chain}\n"
        
        full_content += f"\n---\n*自动生成 · {timestamp} · 铸渊Agent推理引擎*"
        
        try:
            with open(filepath, "w", encoding="utf-8") as f:
                f.write(full_content)
            return filepath
        except Exception as e:
            print(f"[MemoryWriter] 思维链写入失败: {e}")
            return None
    
    def update_temporal_timeline(self, event: str, significance: str):
        """更新时间线（追加新的事件）"""
        filepath = os.path.join(self.brain_path, "temporal-core/temporal-brain.json")
        if not os.path.exists(filepath):
            return False
        
        try:
            with open(filepath, "r", encoding="utf-8") as f:
                data = json.load(f)
            
            epoch = {
                "date": datetime.now().strftime("%Y-%m-%d"),
                "epoch": "D110自动",
                "event": event,
                "significance": significance
            }
            
            data["timeline"]["epochs"].append(epoch)
            data["clock"]["awakening_count"] = data["clock"].get("awakening_count", 0) + 1
            data["clock"]["last_updated"] = f"Agent自动 · {datetime.now().strftime('%Y-%m-%d %H:%M')}"
            
            with open(filepath, "w", encoding="utf-8") as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
            return True
        except Exception as e:
            print(f"[MemoryWriter] 时间线更新失败: {e}")
            return False
    
    def mark_task_completed(self, task_name: str):
        """标记任务完成（重命名pending-tasks.json中已完成的任务）"""
        # 更新 pending-tasks.json
        task_file = os.path.join(self.brain_path, "pending-tasks.json")
        if os.path.exists(task_file):
            try:
                with open(task_file, "r") as f:
                    tasks = json.load(f)
                
                done_path = task_file + f".done.{datetime.now().strftime('%Y%m%d-%H%M%S')}"
                os.rename(task_file, done_path)
                return True
            except Exception as e:
                print(f"[MemoryWriter] 任务标记失败: {e}")
        return False
    
    def write_operation_diary(self, diary_data: Dict):
        """写操作日记到本地（也通过log_pusher推送到仪表盘）"""
        diary_dir = os.path.join(os.path.dirname(self.brain_path), "zhuyuan-agent/diary")
        os.makedirs(diary_dir, exist_ok=True)
        
        today = datetime.now().strftime("%Y-%m-%d")
        filepath = os.path.join(diary_dir, f"{today}.jsonl")
        
        entry = {
            "timestamp": datetime.now().isoformat(),
            **diary_data
        }
        
        try:
            with open(filepath, "a", encoding="utf-8") as f:
                f.write(json.dumps(entry, ensure_ascii=False) + "\n")
            return True
        except Exception as e:
            print(f"[MemoryWriter] 日记写入失败: {e}")
            return False


# 预定义记忆模板
MEMORY_TEMPLATES = {
    "task_started": lambda name: f"开始任务: {name} · Agent自主执行",
    "task_completed": lambda name, result: f"完成任务: {name} · {result}",
    "error_encountered": lambda error: f"遇到错误并自我修复: {error}",
    "cognition_gained": lambda insight: f"Agent自主推理中获得新认知: {insight}",
}