import sys
import json
import importlib.util
import socket
import struct
import traceback
import asyncio
from pathlib import Path

# 尝试导入 langchain 的序列化工具 (如果环境中有)
try:
    from langchain_core.messages import BaseMessage
    from pydantic import BaseModel
except ImportError as e:
    raise ImportError("请安装 langchain_core 和 pydantic 库以支持 LangChain/LangGraph 序列化") from e

try:
    import alayaflow
except ImportError as e:
    raise ImportError("alayaflow没有在worker内安装") from e

try:
    from alayaflow.workflow.workflow_loader import WorkflowLoader
    from alayaflow.common.config import settings
    from alayaflow.execution.langfuse_tracing import get_tracing
except ImportError as e:
    raise e

port = int(sys.argv[1])
wf_dir = sys.argv[2]

s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
    s.connect(('127.0.0.1', port))
except ConnectionRefusedError:
    sys.exit(1)

def default_encoder(obj):
    if hasattr(obj, 'model_dump') and callable(obj.model_dump):
        return obj.model_dump()
    return str(obj)

def send_chunk(data_dict):
    try:
        json_bytes = json.dumps(data_dict, default=default_encoder).encode('utf-8')
        header = struct.pack('!I', len(json_bytes))
        s.sendall(header + json_bytes)
    except Exception as e:
        err_msg = json.dumps({"error": f"Serialization Error: {str(e)}"}).encode('utf-8')
        s.sendall(struct.pack('!I', len(err_msg)) + err_msg)

async def run():
    try:
        input_str = sys.stdin.read()
        lines = input_str.strip().split('\n')
        init_args = json.loads(lines[0]) if len(lines) > 0 and lines[0].strip() else {}
        config = json.loads(lines[1]) if len(lines) > 1 and lines[1].strip() else {}
        input_data = json.loads(lines[2]) if len(lines) > 2 and lines[2].strip() else {}

        print("Input data for workflow:", input_data)

        # 加载工作流元信息和创建器
        # wf_dir 格式: {workflows_dir}/{workflow_id}/{version}
        wf_dir_path = Path(wf_dir)
        workflows_dir = str(wf_dir_path.parent.parent)  # 工作流根目录
        workflow_id = wf_dir_path.parent.name  # 工作流ID
        version = wf_dir_path.name  # 版本号

        workflow_dir = Path(workflows_dir) / workflow_id / version
        
        loader = WorkflowLoader(workflow_dir, workflow_id, version)
        result = loader.load_workflow()
        wf_info = result.info
        creator = result.creator
        graph = creator(init_args)

        tracing = get_tracing(settings)
        langfuse_cb = tracing.build_callback()
        
        # 合并用户提供的 config 和 tracing config
        merged_config = dict(config) if config else {}
        if langfuse_cb:
            merged_config.update(tracing.build_config(input_data, wf_info, langfuse_cb))
        
        async for chunk in graph.astream_events(input_data, config=merged_config, version="v2"):
            send_chunk(chunk)

    except Exception as e:
        error_info = {
            "error": str(e),
            "traceback": traceback.format_exc()
        }
        send_chunk(error_info)
        traceback.print_exc()
        print(f"[Runner Critical] {{error_info}}", file=sys.stderr)

    finally:
        s.close()

asyncio.run(run())
