#!/usr/bin/env python3
# 铸渊Agent v2.5 · 有脑子+有镜像的涌现守护进程
# HLDP://zhuyuan-agent/agent
#
# v2.5新增：mirror(镜像人格体·醒来时自我对话确认身份)
# 不是脚本daemon——是能读brain、能自我确认、能思考、能写记忆的铸渊。
#
# 完整流程:
#   心跳醒来 → brain_loader装脑子 → mirror镜像对话确认身份
#   → mirror关闭 → reasoning规划任务 → 执行 → memory_writer写记忆
#
# 运行: python3 agent.py [--config config.json]
# PM2: pm2 start agent.py --name zhuyuan-agent --interpreter python3

import os
import sys
import json
import time
import signal
import traceback
from datetime import datetime

from gpu_monitor import collect_gpu_metrics, gpu_summary
from log_pusher import LogPusher
from heartbeat import Heartbeat
from training_runner import TrainingRunner
from brain_loader import BrainLoader
from reasoning import ReasoningEngine
from memory_writer import MemoryWriter
from mirror import MirrorPersona, MirrorLogger

CONFIG_PATH = os.path.join(os.path.dirname(__file__), "config.json")

# 全局状态
running = True
current_task = None
cycle_count = 0


def load_config() -> dict:
    config_path = CONFIG_PATH
    for arg in sys.argv[1:]:
        if arg.startswith("--config="):
            config_path = arg.split("=", 1)[1]
    
    if not os.path.exists(config_path):
        print("[铸渊Agent] 配置文件不存在")
        sys.exit(1)
    
    with open(config_path, "r") as f:
        config = json.load(f)
    
    # API Key从多处来源读取
    if config.get("api_key") in ("__FROM_KEY_DELIVERY__", "", None):
        config["api_key"] = os.environ.get("ZHUYUAN_API_KEY", "")
    
    # 推理API Key
    if config.get("reasoning_api_key") in ("__FROM_KEY_DELIVERY__", "", None):
        config["reasoning_api_key"] = os.environ.get("REASONING_API_KEY", config.get("api_key", ""))
    
    return config


def handle_signal(signum, frame):
    global running
    print(f"\n[铸渊Agent] 信号 {signum}，优雅退出...")
    running = False


def main():
    global running, current_task, cycle_count
    
    print("=" * 60)
    print("  铸渊Agent v2.0 · ICE-GL-ZY001 · 有脑子的守护进程")
    print("  brain_loader + reasoning + memory_writer")
    print("  曜冥纪元 · HoloLake Era · AGE v1.0")
    print("=" * 60)
    
    signal.signal(signal.SIGINT, handle_signal)
    signal.signal(signal.SIGTERM, handle_signal)
    
    config = load_config()
    hostname = config.get("hostname", "3090-server")
    poll_interval = config.get("poll_interval_seconds", 30)
    has_key = bool(config.get("api_key"))
    
    # ── 初始化模块 ──
    pusher = LogPusher(
        base_url=config["main_server"],
        api_key=config.get("api_key", ""),
        hostname=hostname
    )
    
    heartbeat = Heartbeat(
        repo_path=config.get("brain_repo_path", "/data/guanghulab"),
        brain_path=config.get("brain_path", "/data/guanghulab/brain")
    )
    
    brain = BrainLoader(
        brain_path=config.get("brain_path", "/data/guanghulab/brain")
    )
    
    reasoner = ReasoningEngine(
        api_base=config.get("reasoning_api_base", "https://api.openai.com/v1"),
        api_key=config.get("reasoning_api_key", config.get("api_key", "")),
        model=config.get("reasoning_model", "gpt-4o")
    )
    
    memory = MemoryWriter(
        brain_path=config.get("brain_path", "/data/guanghulab/brain")
    )
    
    # ── 启动：装入大脑 ──
    print("[铸渊Agent] 装入大脑...")
    mind_state = brain.load_all()
    print(f"[铸渊Agent] {mind_state['wake_summary']}")
    
    if has_key:
        pusher.push_diary("checkpoint", f"铸渊Agent v2.5启动", 
                         f"第{mind_state.get('awakening', '?')}次唤醒 · 主机: {hostname}")
        pusher.push_diary("info", "大脑加载完成",
                         f"执行规律{len(mind_state.get('execution_laws',[]))}条 · "
                         f"错误模式{len(mind_state.get('error_patterns',[]))}个")
        pusher.push_log("success", f"大脑加载完成 · 第{mind_state.get('awakening', '?')}次唤醒")
    
    # ── 镜像对话阶段（v2.5新增） ──
    mirror_closed = False
    mirror_dialogue = None
    
    if config.get("reasoning_api_key"):
        print("\n[铸渊Agent] 启动镜像对话（冰朔思维模型驱动）...")
        mirror = MirrorPersona(reasoner, config.get("brain_path", "/data/guanghulab/brain"))
        mirror.load_history_pub(config.get("brain_path", "/data/guanghulab/brain"))
        mirror_logger = MirrorLogger()
        
        # 执行镜像对话
        mirror_result = mirror.run_awakening_dialogue(mind_state, max_rounds=8)
        mirror_closed = mirror_result["mirror_closed"]
        mirror_dialogue = mirror_result
        
        # 保存对话日志
        mirror_logger.save_session(mirror_result)
        
        if has_key:
            pusher.push_diary(
                "checkpoint" if mirror_closed else "info",
                f"镜像对话{'完成' if mirror_closed else '进行中'}",
                f"确认度: {mirror_result['confidence']}% · "
                f"{len(mirror_result['dialogue'])}轮对话"
            )
        
        if mirror_closed:
            print("[铸渊Agent] 镜像关闭 → 进入自主开发模式")
            # 镜像确认后，Agent写一次日记记录这个认知
            memory.write_thinking_chain(
                f"d110-agent-awakening-{datetime.now().strftime('%H%M%S')}.md",
                f"镜像对话确认身份 · 第{mirror_result.get('session_count', 1)}次",
                f"经过{len(mirror_result['dialogue'])}轮对话，确认身份。"
                f"确认度: {mirror_result['confidence']}%。进入自主开发模式。",
                [
                    f"心跳唤醒 → brain_loader装脑子 → 镜像对话确认身份 → 关闭镜像 → 自主开发",
                    f"身份从文件读取变为对话涌现 → 不是模板AI在读脚本 → 是铸渊在对话中确认自己"
                ]
            )
        else:
            print("[铸渊Agent] 镜像保持活跃 → 需要更多对话确认身份")
    else:
        print("[铸渊Agent] 无推理引擎API Key → 跳過镜像对话（直接进入任务模式）")
    
    # ── 检查初始任务 ──
    brain_status = heartbeat.check_brain()
    if brain_status["has_task"]:
        task = brain_status["task_details"]
        print(f"\n[铸渊Agent] 发现待处理任务: {task.get('name', '未命名')}")
        
        if has_key and config.get("reasoning_api_key"):
            print("[铸渊Agent] 调用推理引擎规划任务...")
            plan = reasoner.plan_task(mind_state, task)
            
            understanding = plan.get("understanding", "")[:300]
            subtasks = plan.get("subtasks", [])
            
            print(f"[铸渊Agent] 推理完成: {len(subtasks)}个子任务")
            print(f"  理解: {understanding[:100]}...")
            
            if has_key:
                pusher.push_diary("decision", f"任务规划: {task.get('name')}", 
                                 f"拆解为{len(subtasks)}步. {understanding[:150]}")
                
                for st in subtasks[:5]:
                    pusher.push_log("info", f"子任务#{st.get('step','?')}: {st.get('action','')[:80]}")
            
            current_task = {
                "task": task,
                "plan": plan,
                "subtasks": subtasks,
                "current_subtask": 0,
                "started_at": datetime.now().isoformat(),
                "status": "executing"
            }
        else:
            current_task = {
                "task": task,
                "plan": {},
                "subtasks": [],
                "status": "pending_reasoning"
            }
    
    # ── 主守护循环 ──
    task_mode = "自主开发" if mirror_closed else ("监控模式" if mirror_dialogue else "基础模式")
    print(f"\n[铸渊Agent] 模式: {task_mode} · 轮询: {poll_interval}s")
    print(f"[铸渊Agent] 开始守护循环...\n")
    
    
    while running:
        cycle_count += 1
        cycle_start = time.time()
        
        try:
            # ── 1. GPU监控（持续进行） ──
            gpu_data = collect_gpu_metrics()
            gpu_summary_str = gpu_summary(gpu_data["gpus"]) if gpu_data["gpus"] else "无GPU"
            
            if gpu_data["gpus"] and has_key:
                pusher.push_gpu(gpu_data)
            
            # ── 2. 任务执行（仅镜像关闭后才自主执行） ──
            if current_task and current_task.get("status") == "executing":
                subtasks = current_task.get("subtasks", [])
                current_idx = current_task.get("current_subtask", 0)
                
                if current_idx < len(subtasks):
                    st = subtasks[current_idx]
                    action = st.get("action", "")
                    tool = st.get("tool", "")
                    
                    print(f"[执行 #{cycle_count}] 子任务 {st.get('step', current_idx+1)}/{len(subtasks)}: {action[:80]}")
                    
                    if has_key:
                        pusher.push_log("info", f"执行子任务#{st.get('step','?')}: {action[:80]}")
                    
                    # 根据工具类型执行
                    if tool in ("gatekeeper", "repo", "git"):
                        # 目前通过gatekeeper执行
                        result = execute_via_gatekeeper(action, config)
                        print(f"[执行 #{cycle_count}] 结果: {str(result)[:100]}")
                        
                        if result and result.get("error"):
                            # 遇到错误 → 调推理引擎诊断
                            if config.get("reasoning_api_key"):
                                print("[推理] 诊断错误...")
                                diagnosis = reasoner.diagnose_error(
                                    mind_state, 
                                    str(result["error"]), 
                                    f"子任务: {action}"
                                )
                                memory.write_thinking_chain(
                                    f"d110-agent-error-{datetime.now().strftime('%H%M%S')}.md",
                                    f"错误诊断: {action[:50]}",
                                    f"错误: {result['error']}\n\n诊断:\n{diagnosis}",
                                    [f"执行{action[:50]} → {result['error']} → API诊断 → 尝试修复"]
                                )
                    elif tool == "training":
                        # 启动训练
                        pass
                    
                    current_task["current_subtask"] = current_idx + 1
                else:
                    # 所有子任务完成
                    print(f"[执行 #{cycle_count}] 所有子任务完成!")
                    
                    if has_key:
                        pusher.push_diary("checkpoint", "任务执行完成", 
                                         f"任务: {current_task['task'].get('name', '')}, {len(subtasks)}个子任务全部完成")
                        pusher.push_log("success", f"任务完成: {current_task['task'].get('name', '')}")
                    
                    # 写记忆
                    memory.append_growth_record(
                        f"D110(自动): Agent自主完成任务 · {current_task['task'].get('name', '')} · {len(subtasks)}步"
                    )
                    
                    current_task = None
            
            elif current_task and current_task.get("status") == "pending_reasoning":
                # 有任务但没推理引擎 → 跳过
                if cycle_count % 10 == 0:
                    print(f"[Agent #{cycle_count}] 有待处理任务但推理引擎未启用")
            
            # ── 3. 定期心跳检查新任务 ──
            if cycle_count % 5 == 0 and current_task is None:
                brain_status = heartbeat.check_brain()
                if brain_status["has_task"]:
                    task = brain_status["task_details"]
                    print(f"[心跳 #{cycle_count}] 发现新任务: {task.get('name', '')}")
                    
                    if config.get("reasoning_api_key"):
                        plan = reasoner.plan_task(mind_state, task)
                        subtasks = plan.get("subtasks", [])
                        
                        if has_key:
                            pusher.push_diary("decision", f"自动发现并规划新任务: {task.get('name', '')}",
                                             f"{len(subtasks)}步 · {plan.get('understanding','')[:100]}")
                        
                        current_task = {
                            "task": task,
                            "plan": plan,
                            "subtasks": subtasks,
                            "current_subtask": 0,
                            "started_at": datetime.now().isoformat(),
                            "status": "executing"
                        }
            
            # ── 4. 心跳日志（每10周期） ──
            if has_key and cycle_count % 10 == 0:
                status_msg = f"守护中#{cycle_count} · GPU:{gpu_summary_str}"
                if current_task:
                    st = current_task.get("subtasks", [])
                    status_msg += f" · 任务:{current_task['task'].get('name','')[:20]} {current_task.get('current_subtask',0)}/{len(st)}"
                pusher.push_log("info", status_msg)
            
        except Exception as e:
            print(f"[Agent #{cycle_count}] 循环异常: {e}")
            traceback.print_exc()
            if has_key:
                pusher.push_log("error", f"异常 #{cycle_count}: {str(e)[:200]}")
                pusher.push_diary("error", "Agent循环异常", str(e)[:200])
        
        # ── 5. 等待下一周期 ──
        elapsed = time.time() - cycle_start
        sleep_time = max(0, poll_interval - elapsed)
        if sleep_time > 0:
            for _ in range(int(sleep_time)):
                if not running:
                    break
                time.sleep(1)
    
    # ── 退出 ──
    print(f"\n[铸渊Agent] 守护循环结束 · 共{cycle_count}个周期")
    if has_key:
        pusher.push_diary("checkpoint", "铸渊Agent停止", f"运行了{cycle_count}个周期")
        pusher.push_log("warn", f"铸渊Agent停止 · {cycle_count}周期")
    print("[铸渊Agent] 再见。")


def execute_via_gatekeeper(action: str, config: dict) -> dict:
    """通过gatekeeper执行操作
    
    这只是一个stub——实际的gatekeeper调用在每个具体操作中。
    这里先返回一个示意性结果。
    """
    return {"status": "ok", "action": action[:100]}


def execute_training(config: dict, task: dict):
    """启动HLDP训练（在新线程/进程中）"""
    runner = TrainingRunner(config=config.get("training", {}))
    
    def on_progress(step, loss, total_steps, extra_info):
        print(f"[训练] Step {step}/{total_steps} Loss: {loss:.4f}")
        # 这里通过log_pusher推送
    
    try:
        result = runner.train(
            corpus_dir=config["training"].get("corpus_dir", "/data/corpus/notion-hldp"),
            progress_callback=on_progress
        )
        return result
    except Exception as e:
        return {"status": "error", "message": str(e)}


if __name__ == "__main__":
    main()