import json
import sqlite3
import time
from typing import Dict, Any, List
from pydantic import BaseModel

class Checkpoint(BaseModel):
    task_id: str
    step: int
    state: Dict[str, Any]
    messages: List[Dict[str, Any]]
    updated_at: float

class ProductionHarnessStore:
    def __init__(self, db_path: str = "agent_state.db"):
        self.conn = sqlite3.connect(db_path)
        self._init_db()

    def _init_db(self):
        with self.conn:
            self.conn.execute("""
                CREATE TABLE IF NOT EXISTS checkpoints (
                    task_id TEXT PRIMARY KEY,
                    step INTEGER,
                    payload TEXT,
                    updated_at REAL
                )
            """)

    def save_checkpoint(self, checkpoint: Checkpoint):
        with self.conn:
            self.conn.execute("""
                INSERT OR REPLACE INTO checkpoints (task_id, step, payload, updated_at)
                VALUES (?, ?, ?, ?)
            """, (
                checkpoint.task_id,
                checkpoint.step,
                checkpoint.model_dump_json(),
                time.time()
            ))

    def load_checkpoint(self, task_id: str) -> Checkpoint | None:
        cursor = self.conn.cursor()
        cursor.execute("SELECT payload FROM checkpoints WHERE task_id = ?", (task_id,))
        row = cursor.fetchone()
        if row:
            return Checkpoint.model_validate_json(row[0])
        return None

def run_resumable_task(task_id: str, initial_goal: str, store: ProductionHarnessStore):
    # 尝试恢复已有的断点
    checkpoint = store.load_checkpoint(task_id)
    
    if checkpoint:
        print(f"[恢复任务 {task_id}] 从第 {checkpoint.step} 步恢复...")
        current_step = checkpoint.step
        state = checkpoint.state
        messages = checkpoint.messages
    else:
        print(f"[新建任务 {task_id}] 初始化状态机...")
        current_step = 0
        state = {"goal": initial_goal, "status": "IN_PROGRESS", "completed_work": []}
        messages = [{"role": "user", "content": initial_goal}]

    max_steps = 5
    while current_step < max_steps:
        current_step += 1
        print(f"\n--- 执行第 {current_step} 步 ---")
        
        # 模拟模型调用与工具操作
        time.sleep(0.5)
        step_result = f"完成阶段子任务 {current_step}"
        
        # 更新状态
        state["completed_work"].append(step_result)
        messages.append({"role": "assistant", "content": step_result})
        
        # 保存持久化检查点（在发生崩溃前落盘）
        store.save_checkpoint(Checkpoint(
            task_id=task_id,
            step=current_step,
            state=state,
            messages=messages,
            updated_at=time.time()
        ))
        print(f"[持久化] 第 {current_step} 步检查点保存成功")

        if current_step >= 3:
            # 达成业务目标提前终止
            state["status"] = "SUCCESS"
            print("[任务完成] 全部业务指标达标。")
            break