import importlib
import os
import subprocess
import tempfile
import json
import asyncio
from typing import Dict, Any


class NodeExecutor:
	"""节点执行器 - 支持多语言节点执行"""

	def __init__(self):
		self.nodes_base_path = "ai/nodes"
	
	async def execute_node(self, node_type: str, flow_id: str, context: Dict[str, Any], 
	                     inputs: Dict[str, Any], config: Dict[str, Any], session_cookie: str = None) -> Dict[str, Any]:
		"""执行节点"""
		try:
			# 检测节点实现文件
			implementation_file = self._detect_implementation(node_type)
			
			# 根据文件扩展名选择执行器
			if implementation_file.endswith('.py'):
				return await self._execute_python_node(node_type, flow_id, context, inputs, config, session_cookie)
			elif implementation_file.endswith('.go'):
				return await self._execute_go_node(node_type, context, inputs, config)
			elif implementation_file.endswith('.rs'):
				return await self._execute_rust_node(node_type, context, inputs, config)
			else:
				return {"success": False, "error": f"不支持的节点实现语言: {node_type}"}
				
		except Exception as e:
			return {"success": False, "error": str(e)}
	
	def _detect_implementation(self, node_type: str) -> str:
		"""检测节点实现文件"""
		# 获取当前文件所在目录的绝对路径
		current_dir = os.path.dirname(os.path.abspath(__file__))
		jingrow_root = os.path.dirname(os.path.dirname(os.path.dirname(current_dir)))
		node_path = os.path.join(jingrow_root, self.nodes_base_path, node_type)
		
		# 按优先级检查实现文件
		for ext in ['.py', '.go', '.rs', '.js']:
			impl_file = os.path.join(node_path, f"{node_type}{ext}")
			if os.path.exists(impl_file):
				return impl_file
		
		raise FileNotFoundError(f"节点 {node_type} 没有找到实现文件，查找路径: {node_path}")
	
	async def _execute_python_node(self, node_type: str, flow_id: str, context: Dict[str, Any], 
	                            inputs: Dict[str, Any], config: Dict[str, Any], session_cookie: str = None) -> Dict[str, Any]:
		"""执行Python节点"""
		try:
			# 动态导入Python模块
			module_path = f"jingrow.ai.nodes.{node_type}.{node_type}"
			module = importlib.import_module(module_path)
			
			# 设置session_cookie和flow_id到context
			if session_cookie:
				context['session_cookie'] = session_cookie
			context['flow_id'] = flow_id
			
			if hasattr(module, 'execute'):
				# 在异步环境中执行同步函数，传递包含session_cookie的context
				loop = asyncio.get_event_loop()
				result = await loop.run_in_executor(
					None, 
					self._execute_with_context,
					module.execute,
					context, inputs, config
				)
				return result
			else:
				return {"success": False, "error": f"Python节点 {node_type} 没有 execute 函数"}
				
		except Exception as e:
			return {"success": False, "error": f"Python执行失败: {str(e)}"}
	
	def _execute_with_context(self, execute_func, context, inputs, config):
		"""在新线程中执行节点函数并设置context"""
		from jingrow.utils.auth import set_context
		set_context(context)
		return execute_func(context, inputs, config)
	
	async def _execute_go_node(self, node_type: str, context: Dict[str, Any], 
	                        inputs: Dict[str, Any], config: Dict[str, Any]) -> Dict[str, Any]:
		"""执行Go节点"""
		try:
			# 读取Go代码
			go_file = f"{self.nodes_base_path}/{node_type}/{node_type}.go"
			with open(go_file, 'r', encoding='utf-8') as f:
				go_code = f.read()
			
			# 创建临时Go文件
			with tempfile.NamedTemporaryFile(mode='w', suffix='.go', delete=False, encoding='utf-8') as f:
				f.write(go_code)
				temp_file = f.name
			
			try:
				# 准备输入数据
				input_data = json.dumps({
					'context': context,
					'inputs': inputs,
					'config': config
				})
				
				# 在异步环境中执行Go程序
				loop = asyncio.get_event_loop()
				result = await loop.run_in_executor(
					None,
					self._run_go_program,
					temp_file,
					input_data
				)
				
				return result
				
			finally:
				# 清理临时文件
				if os.path.exists(temp_file):
					os.unlink(temp_file)
				
		except Exception as e:
			return {"success": False, "error": f"Go执行失败: {str(e)}"}
	
	async def _execute_rust_node(self, node_type: str, context: Dict[str, Any], 
	                           inputs: Dict[str, Any], config: Dict[str, Any]) -> Dict[str, Any]:
		"""执行Rust节点"""
		try:
			# 读取Rust代码
			rust_file = f"{self.nodes_base_path}/{node_type}/{node_type}.rs"
			with open(rust_file, 'r', encoding='utf-8') as f:
				rust_code = f.read()
			
			# 创建临时Rust文件
			with tempfile.NamedTemporaryFile(mode='w', suffix='.rs', delete=False, encoding='utf-8') as f:
				f.write(rust_code)
				temp_file = f.name
			
			try:
				# 准备输入数据
				input_data = json.dumps({
					'context': context,
					'inputs': inputs,
					'config': config
				})
				
				# 在异步环境中执行Rust程序
				loop = asyncio.get_event_loop()
				result = await loop.run_in_executor(
					None,
					self._run_rust_program,
					temp_file,
					input_data
				)
				
				return result
				
			finally:
				# 清理临时文件
				if os.path.exists(temp_file):
					os.unlink(temp_file)
					
		except Exception as e:
			return {"success": False, "error": f"Rust执行失败: {str(e)}"}
	
	def _run_go_program(self, temp_file: str, input_data: str) -> Dict[str, Any]:
		"""运行Go程序"""
		try:
			# 编译并执行Go程序
			result = subprocess.run(
				['go', 'run', temp_file],
				input=input_data,
				capture_output=True,
				text=True,
				timeout=30
			)
			
			if result.returncode == 0:
				try:
					# 尝试解析JSON输出
					output_data = json.loads(result.stdout.strip())
					return output_data
				except json.JSONDecodeError:
					return {
						"success": True,
						"output": result.stdout,
						"raw_output": True
					}
			else:
				return {
					"success": False,
					"error": result.stderr,
					"returncode": result.returncode
				}
				
		except subprocess.TimeoutExpired:
			return {"success": False, "error": "Go程序执行超时"}
		except Exception as e:
			return {"success": False, "error": str(e)}
	
	def _run_rust_program(self, temp_file: str, input_data: str) -> Dict[str, Any]:
		"""运行Rust程序"""
		try:
			# 编译Rust程序
			compile_result = subprocess.run(
				['rustc', temp_file, '-o', '/tmp/rust_temp'],
				capture_output=True,
				text=True,
				timeout=30
			)
			
			if compile_result.returncode != 0:
				return {
					"success": False,
					"error": f"Rust编译失败: {compile_result.stderr}"
				}
			
			try:
				# 执行编译后的程序
				exec_result = subprocess.run(
					['/tmp/rust_temp'],
					input=input_data,
					capture_output=True,
					text=True,
					timeout=30
				)
				
				if exec_result.returncode == 0:
					try:
						# 尝试解析JSON输出
						output_data = json.loads(exec_result.stdout.strip())
						return output_data
					except json.JSONDecodeError:
						return {
							"success": True,
							"output": exec_result.stdout,
							"raw_output": True
						}
				else:
					return {
						"success": False,
						"error": exec_result.stderr,
						"returncode": exec_result.returncode
					}
					
			finally:
				# 清理编译后的文件
				if os.path.exists('/tmp/rust_temp'):
					os.unlink('/tmp/rust_temp')
					
		except subprocess.TimeoutExpired:
			return {"success": False, "error": "Rust程序执行超时"}
		except Exception as e:
			return {"success": False, "error": str(e)}


