import os
import sys

from dotenv import load_dotenv
from alayaflow.api import Flow

# Load environment variables from .env file
load_dotenv()


def init_flow():
    """Initialize flow api singleton"""

    # Flow() is a singleton instance
    # So it can be "constructed" in multiple places
    # But there is only one instance
    flow = Flow()

    flow.init({
        "alayahub_url": "http://your-alayahub-url",
    })

    # api_key = os.getenv("DEEPSEEK_API_KEY")
    # if not api_key:
    #     raise ValueError("请设置环境变量 DEEPSEEK_API_KEY")

    flow.register_models([
        # For test
        # {
        #     # Local used fields
        #     "name": "DeepSeek Chat",
        #     "model_id": "alayajet",
        #     "provider_name": "DeepSeek",
        #     # Connection credentials
        #     "model_name": "deepseek-chat",
        #     "base_url": "https://api.deepseek.com/v1",
        #     "api_key": api_key,
        # }
        {
            # Local used fields
            "name": "AlayaJet Chat",
            "model_id": "alayajet",
            "provider_name": "Alaya.AI",
            # Connection credentials
            "model_name": "<your-model-name>",
            "base_url": "http://localhost:12345",
            "api_key": "",
        }
    ])


def init_workflows():
    flow = Flow()

    workflow_id = 'mem_chat'
    workflow_version = '1.0.0'
    init_args = {
        "alayamem_url": "http://100.64.0.5:5556",
    }
    print(f"\n加载工作流: {workflow_id} v{workflow_version}")
    try:
        flow.load_workflow(workflow_id, workflow_version, init_args)
    except Exception as e:
        import traceback
        traceback.print_exc()


def run_single_turn(flow, workflow_id, workflow_version, question, context):
    """执行单轮对话"""
    inputs = {
        "messages": [{"role": "user", "content": question}]
    }

    def handle_chat_model_stream(event):
        content = event["data"]["chunk"]['content']
        if content:
            print(content, end="", flush=True)

    def handle_error(event):
        sys.stderr.write(f"\n[Error Event Detected]\n")
        sys.stderr.write("--- Error message ---\n")
        sys.stderr.write(f"{event.get('error', 'Unknown error')}\n")
        if 'traceback' in event:
            sys.stderr.write("\n--- Stack Trace ---\n")
            sys.stderr.write(f"{event['traceback']}\n")
        # 打印完整的 event 信息以便调试
        sys.stderr.write("\n--- Complete Event ---\n")
        sys.stderr.write(f"{event}\n")

    for event in flow.exec_workflow(workflow_id, workflow_version, inputs, context):
        if "error" in event:
            handle_error(event)
            return False
        kind = event["event"]
        if kind == "on_chat_model_stream":
            handle_chat_model_stream(event)
    print()
    return True


def run_workflow():
    """运行多轮对话测试"""
    # 使用已初始化的 Flow 单例
    flow = Flow()

    workflow_id = 'mem_chat'
    workflow_version = '1.0.0'
    context = {
        "user_id": "kekeke",
        "session_id": "test_session_multi_turn",
        "chat_model_id": "alayajet",
        "load_kv_cache_path": "<your-kv-cache-path>",
    }

    # 定义多轮对话问题
    questions = [
        "你是谁",
        "3句话介绍一下transformers",
        "我问了你几个问题",
        "你觉得我是做什么工作的"
    ]

    print("\n" + "="*80)
    print("开始多轮对话测试 - 测试 AlayaMem 记忆功能")
    print("="*80)

    for i, question in enumerate(questions, 1):
        print(f"\n{'='*80}")
        print(f"第 {i} 轮对话")
        print(f"{'='*80}")
        print(f"\n用户: {question}")
        print(f"助手: ", end="")

        success = run_single_turn(flow, workflow_id, workflow_version, question, context)
        if not success:
            print(f"\n第 {i} 轮对话失败，停止测试")
            break

        print(f"\n{'-'*80}")

    print(f"\n{'='*80}")
    print("多轮对话测试完成")
    print(f"{'='*80}\n")


def main():
    init_flow()
    init_workflows()
    run_workflow()


if __name__ == "__main__":
    main()