import os
import sys

from alayaflow.api import Flow
from dotenv import load_dotenv
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([
        {
            # Local used fields
            "name": "DeepSeek Chat",
            "model_id": "deepseek-chat",
            "provider_name": "DeepSeek",
            # Connection credentials
            "model_name": "deepseek-chat",
            "base_url": "https://api.deepseek.com/v1",
            "api_key": api_key,
        }
    ])


def init_workflows():
    flow = Flow()

    workflow_id = 'simple_chat'
    workflow_version = '1.0.0'
    init_args = {
        "alayamem_url": "http://100.64.0.5:5556",
    }
    flow.load_workflow(workflow_id, workflow_version, init_args)


def run_workflow():
    flow = Flow()

    workflow_id = 'simple_chat'
    workflow_version = '1.0.0'
    inputs = {
        "messages": [{"role": "user", "content": "你是谁"}]
    }
    context = {
        "user_id": "test_user",
        "session_id": "test_session",
        "chat_model_id": "deepseek-chat",
    }

    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['error']}\n")
        sys.stderr.write("\n--- Stack Trace ---\n")
        sys.stderr.write(f"{event['traceback']}\n")

    print("Start executing workflow...")
    for event in flow.exec_workflow(workflow_id, workflow_version, inputs, context):
        if "error" in event:
            handle_error(event)
            break
        kind = event["event"]
        if kind == "on_chat_model_stream":
            handle_chat_model_stream(event)
    print()


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


if __name__ == "__main__":
    main()