from flask import request, jsonify
import asyncio
import uuid
from graph.test_graph_3 import run_with_persistence, get_checkpoint_history, resume_from_checkpoint

def run_async(coro):
    """运行异步函数的辅助函数"""
    try:
        # 尝试使用现有的事件循环
        loop = asyncio.get_running_loop()
    except RuntimeError:
        # 如果没有运行中的事件循环，则创建一个新的
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
    
    return loop.run_until_complete(coro)

def run_langgraph():
    """启动一个新的langgraph任务"""
    try:
        data = request.get_json()
        user_input = data.get('user_input', '')
        thread_id = data.get('thread_id', str(uuid.uuid4()))
        
        if not user_input:
            return jsonify({'error': 'user_input is required'}), 400
        
        # 运行异步函数
        output, thread_id = run_async(run_with_persistence(user_input, thread_id))
        
        return jsonify({
            'status': 'success',
            'thread_id': thread_id,
            'output': output
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

def get_task_status(thread_id):
    """查询任务状态和历史"""
    try:
        # 获取检查点历史
        # 注意：这里需要修改get_checkpoint_history以返回数据而不是打印
        # history = run_async(get_checkpoint_history(thread_id))
        
        return jsonify({
            'status': 'success',
            'thread_id': thread_id,
            'message': 'Task status endpoint'
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500

def resume_task(thread_id):
    """从检查点恢复任务"""
    try:
        data = request.get_json()
        checkpoint_id = data.get('checkpoint_id')
        
        # 恢复检查点状态
        restored_state = resume_from_checkpoint(thread_id, checkpoint_id)
        
        if restored_state:
            return jsonify({
                'status': 'success',
                'restored_state': restored_state
            })
        else:
            return jsonify({'error': 'Failed to restore checkpoint'}), 404
    except Exception as e:
        return jsonify({'error': str(e)}), 500

def visualize_graph(thread_id):
    """可视化图结构"""
    try:
        # 这里可以返回图的可视化信息
        # 为了简化，我们只返回基本信息
        return jsonify({
            'status': 'success',
            'thread_id': thread_id,
            'message': 'Graph visualization endpoint'
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500