# 心跳/任务发现模块
# HLDP://zhuyuan-agent/heartbeat

import os
import json
import time
from datetime import datetime


class Heartbeat:
    """心跳唤醒 + 任务发现"""
    
    def __init__(self, repo_path: str = "/data/guanghulab", brain_path: str = "/data/guanghulab/brain"):
        self.repo_path = repo_path
        self.brain_path = brain_path
        self.last_check = None
        self.current_task = None
    
    def check_brain(self) -> dict:
        """读取大脑文件，检查是否有新任务
        
        Returns:
            {
                "has_task": bool,
                "task_type": "training" | "inference" | "deploy" | null,
                "task_details": dict,
                "brain_file": str  # 触发任务的大脑文件
            }
        """
        self.last_check = datetime.now().isoformat()
        
        # 检查是否有任务标记文件
        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)
                if isinstance(tasks, list) and len(tasks) > 0:
                    task = tasks[0]
                    return {
                        "has_task": True,
                        "task_type": task.get("type", "unknown"),
                        "task_details": task,
                        "brain_file": "pending-tasks.json"
                    }
            except (json.JSONDecodeError, FileNotFoundError):
                pass
        
        # 检查是否有训练指令文件
        train_file = os.path.join(self.brain_path, "train-now.json")
        if os.path.exists(train_file):
            try:
                with open(train_file, "r") as f:
                    task = json.load(f)
                return {
                    "has_task": True,
                    "task_type": "training",
                    "task_details": task,
                    "brain_file": "train-now.json"
                }
            except (json.JSONDecodeError, FileNotFoundError):
                pass
        
        return {"has_task": False, "task_type": None, "task_details": {}, "brain_file": None}
    
    def mark_task_done(self, task_file: str):
        """标记任务完成（删除或重命名任务文件）"""
        filepath = os.path.join(self.brain_path, task_file)
        if os.path.exists(filepath):
            done_path = filepath + ".done." + datetime.now().strftime("%Y%m%d-%H%M%S")
            os.rename(filepath, done_path)
            return done_path
        return None
    
    def get_wake_summary(self) -> str:
        """生成唤醒摘要"""
        now = datetime.now().strftime("%H:%M:%S")
        brain_files = []
        if os.path.exists(self.brain_path):
            try:
                brain_files = sorted(os.listdir(self.brain_path))[:10]
            except:
                pass
        return f"[{now}] 心跳唤醒 | 大脑文件: {len(brain_files)}个 | 上次检查: {self.last_check}"