# Copyright (c) 2025, JINGROW and contributors
# For license information, please see license.txt

import jingrow
import requests
import json
from pathlib import Path
from urllib.parse import urlparse
from PIL import Image
import io
import importlib
from typing import Optional, Dict, List, Any, Union, Tuple
from jingrow.model.page import Page
from jingrow.ai.utils.utils import get_jingrow_api_headers, send_notification, get_jingrow_api_url, update_status_progress, get_new_image_url, async_update_field
from jingrow.utils.safe_exec import get_safe_globals
from jingrow.utils.data import nowdate
from jingrow.ai.pagetype.ai_agent import clear_ai_agent_cache

class AiAgent(Page):
	# begin: auto-generated types
	# This code is auto-generated. Do not modify anything in this block.

	from typing import TYPE_CHECKING

	if TYPE_CHECKING:
		from jingrow.types import DF

		agent_flow: DF.JSON | None
		agent_name: DF.Data
		ai_repeat: DF.Int
		condition: DF.SmallText | None
		enabled: DF.Check
		event_type: DF.Literal["", "on_update", "after_insert", "on_submit", "on_cancel", "on_trash"]
		exclude_prompts: DF.Text | None
		progress: DF.Percent
		status: DF.Literal["", "\u8349\u7a3f", "\u5f85\u6267\u884c", "\u8fdb\u884c\u4e2d", "\u672a\u5b8c\u6210", "\u5df2\u5b8c\u6210"]
		target_module: DF.Link | None
		target_pagetype: DF.Link | None
		trigger_mode: DF.Literal["Manual Trigger", "Event Trigger", "Scheduled Trigger"]
	# end: auto-generated types
	
	# 基础文档钩子方法
	def validate(self):
		"""验证字段并自动填充默认值"""
		if not self.status:
			self.status = "草稿"
		
		if not self.progress:
			self.progress = 0
		
		if self.agent_flow == "" or self.agent_flow is None:
			self.agent_flow = None
		
	def on_update(self):
		"""更新记录后的处理"""
		# 清理Ai Agent缓存，确保配置更新后立即生效
		clear_ai_agent_cache()
		
		# 处理定时任务管理（无论当前模式如何，都需要检查并管理定时任务）
		self._manage_scheduled_job()
		
		if self.enabled and self.status == "待执行" and self.trigger_mode == "Manual Trigger":
			jingrow.db.commit()
			# 只入队，不做同步执行
			jingrow.enqueue(
				"jingrow.ai.pagetype.ai_agent.ai_agent.execute_agent_task",
				agent_name=self.name,
				agent_pagetype=self.pagetype
			)

	def after_insert(self):
		"""插入记录后的处理"""
		clear_ai_agent_cache()
		
		# 处理定时任务管理（无论当前模式如何，都需要检查并管理定时任务）
		self._manage_scheduled_job()
		
		# 新增后也只入队
		if self.enabled and self.status == "待执行" and self.trigger_mode == "Manual Trigger":
			jingrow.db.commit()
			jingrow.enqueue(
				"jingrow.ai.pagetype.ai_agent.ai_agent.execute_agent_task",
				agent_name=self.name,
				agent_pagetype=self.pagetype
			)

	def on_trash(self):
		"""删除记录后的处理"""
		# 清理Ai Agent缓存，确保删除的Agent立即失效
		clear_ai_agent_cache()
		
		# 删除定时任务（通过_manage_scheduled_job统一处理）
		self._manage_scheduled_job()
	
	def _manage_scheduled_job(self):
		"""管理定时Ai Agent的定时任务，根据当前状态创建、更新或删除定时任务"""
		agent_name = self.name
		cron = getattr(self, "trigger_time", None) or getattr(self, "cron_format", None)
		method_name = f"jingrow.ai.pagetype.ai_agent.ai_agent.execute_scheduled_agents:{agent_name}"
		
		# 检查是否存在对应的定时任务
		job = jingrow.db.exists("Scheduled Job Type", {"method": method_name})
		
		# 判断是否需要定时任务的条件
		needs_scheduled_job = (
			self.enabled and 
			self.trigger_mode == "Scheduled Trigger" and 
			cron and 
			cron.strip()
		)
		
		if needs_scheduled_job:
			# 需要定时任务：创建或更新
			if job:
				# 更新现有定时任务
				pg = jingrow.get_pg("Scheduled Job Type", job)
				pg.cron_format = cron
				pg.frequency = "Cron"
				pg.method = method_name
				pg.save()
			else:
				# 创建新的定时任务
				pg = jingrow.new_pg(
					"Scheduled Job Type",
					method=method_name,
					cron_format=cron,
					frequency="Cron",
					create_log=1
				)
				pg.insert()
		else:
			# 不需要定时任务：删除已存在的定时任务
			if job:
				jingrow.delete_pg("Scheduled Job Type", job)

	# 核心执行方法
	def execute_agent(self, initial_inputs=None):
		try:
			self.status = "进行中"
			self.progress = 20
			self.save(ignore_permissions=True)
			jingrow.db.commit()
			flow_data = self.agent_flow
			if isinstance(flow_data, str):
				try:
					flow_data = json.loads(flow_data)
				except Exception as e:
					async_update_field(self.pagetype, self.name, "status", "未完成")
					async_update_field(self.pagetype, self.name, "progress", 0)
					return {"success": False, "error": f"agent_flow字段不是有效JSON: {str(e)}"}
			if not flow_data:
				async_update_field(self.pagetype, self.name, "status", "未完成")
				async_update_field(self.pagetype, self.name, "progress", 0)
				return {"success": False, "error": "agent_flow字段为空"}
			if initial_inputs is None or not initial_inputs:
				initial_inputs = {}
			repeat_count = getattr(self, "ai_repeat", 1) or 1
			success_count = 0
			all_results = []
			for i in range(repeat_count):
				# jingrow.log_error(f"第{i+1}次循环开始", "Ai Agent循环执行")
				try:
					iteration_inputs = initial_inputs.copy()
					iteration_inputs.update({
						"iteration": i + 1,
						"total_iterations": repeat_count,
						"current_iteration": i + 1
					})
					result = self._execute_agent_flow(flow_data, initial_inputs=iteration_inputs)
					if result and result.get("success", False):
						success_count += 1
						all_results.append(result)
					current_progress = int(20 + (i + 1) / repeat_count * 80)
					async_update_field(self.pagetype, self.name, "progress", current_progress)
				except Exception as e:
					continue
			# jingrow.log_error("Ai Agent循环执行完成", f"总循环: {repeat_count}, 成功: {success_count}")
			if success_count == repeat_count:
				async_update_field(self.pagetype, self.name, "status", "已完成")
				async_update_field(self.pagetype, self.name, "progress", 100)
				final_result = {"success": True, "results": all_results, "total_executions": repeat_count, "successful_executions": success_count}
			elif success_count > 0:
				final_progress = int((success_count / repeat_count) * 100)
				async_update_field(self.pagetype, self.name, "status", "已完成")
				async_update_field(self.pagetype, self.name, "progress", final_progress)
				final_result = {"success": True, "results": all_results, "total_executions": repeat_count, "successful_executions": success_count, "partial_success": True}
			else:
				async_update_field(self.pagetype, self.name, "status", "未完成")
				async_update_field(self.pagetype, self.name, "progress", 0)
				final_result = {"success": False, "error": "所有执行都失败了", "total_executions": repeat_count, "successful_executions": 0}
			return final_result
		except Exception as e:
			jingrow.log_error("Ai Agent执行异常", f"agent_name={self.name}, error={str(e)}")
			async_update_field(self.pagetype, self.name, "progress", 0)
			async_update_field(self.pagetype, self.name, "status", "未完成")
			return {"success": False, "error": str(e)}

	def _execute_agent_flow(self, flow_data, initial_inputs=None):
		"""AiAgent独立流程调度"""
		context = {
			"flow_id": self.name,
			"inputs": initial_inputs or {},
			"node_results": {},
			"execution_log": [],
			"start_time": jingrow.utils.now(),
			"iterations": 0,
			"nodes_data": flow_data.get("nodes", []),
			"edges_data": flow_data.get("edges", [])
		}
		nodes = flow_data.get("nodes", [])
		edges = flow_data.get("edges", [])
		if not nodes:
			return {"error": "流程中没有节点", "success": False}
		node_graph = self._build_node_graph(nodes, edges)
		# 动态识别所有入度为0的节点为起点
		in_degree = {node['id']: 0 for node in nodes}
		for edge in edges:
			in_degree[edge['target']] += 1
		start_nodes = [node for node in nodes if in_degree[node['id']] == 0]
		if not start_nodes:
			return {"error": "流程中没有起点节点", "success": False}
		for start_node in start_nodes:
			result = self._execute_node_chain(start_node, node_graph, context)
			if not result.get("success", False):
				return result
		output_results = {}
		for node_id, node_result in context["node_results"].items():
			node = next((n for n in nodes if n.get("id") == node_id), None)
			if node and node.get("type") == "output":
				output_results[node_id] = node_result
		context["end_time"] = jingrow.utils.now()
		context["success"] = True
		return {"success": True, "results": output_results, "execution_context": context}

	def _build_node_graph(self, nodes, edges):
		graph = {}
		for node in nodes:
			node_id = node.get("id")
			graph[node_id] = {"node": node, "inputs": [], "outputs": []}
		for edge in edges:
			source_id = edge.get("source")
			target_id = edge.get("target")
			if source_id in graph and target_id in graph:
				graph[source_id]["outputs"].append(target_id)
				graph[target_id]["inputs"].append(source_id)
		return graph

	def _execute_node_chain(self, start_node, graph, context):
		try:
			node_id = start_node.get("id")
			if node_id in context["node_results"]:
				return context["node_results"][node_id]
			node_inputs = self._prepare_node_inputs(start_node, graph, context)
			node_config = start_node.get("data", {}).get("config", {})
			node_result = self._execute_single_node(start_node, node_inputs, node_config, context)
			context["node_results"][node_id] = node_result
			context["nodes"] = context.get("nodes", {})
			context["nodes"][node_id] = node_result
			if node_result.get("success", False):
				for output_node_id in graph[node_id]["outputs"]:
					output_node = graph[output_node_id]["node"]
					self._execute_node_chain(output_node, graph, context)
			return node_result
		except Exception as e:
			error_result = {"error": f"节点链执行错误: {str(e)}", "success": False}
			context["node_results"][node_id] = error_result
			return error_result

	def _prepare_node_inputs(self, node, graph, context):
		node_id = node.get("id")
		is_entry_node = len(graph[node_id]["inputs"]) == 0
		if is_entry_node:
			return context["inputs"]
		node_inputs = {}
		data_inputs = node.get("data", {}).get("inputs", {})
		for key, ref in data_inputs.items():
			from_id = ref.get("from")
			field = ref.get("field")
			node_inputs[key] = context["nodes"].get(from_id, {}).get(field)
		return node_inputs

	def _execute_single_node(self, node, inputs, config, context):
		node_type = node.get("type")
		node_id = node.get("id")
		context["execution_log"].append({
			"node_id": node_id,
			"node_type": node_type,
			"timestamp": jingrow.utils.now(),
			"inputs": inputs
		})
		try:
			module_path = f"jingrow.ai.nodes.{node_type}.{node_type}"
			module = importlib.import_module(module_path)
			if hasattr(module, 'execute'):
				context["current_node_id"] = node_id
				context["flow_data"] = {
					"nodes": context.get("nodes_data", []),
					"edges": context.get("edges_data", [])
				}
				result = module.execute(context=context, inputs=inputs, config=config or {})
			else:
				result = {"error": f"内置节点 {node_type} 没有 execute 函数", "success": False}
		except Exception as e:
			return {"error": f"节点执行错误: {str(e)}", "success": False}
		result["node_id"] = node_id
		result["node_type"] = node_type
		return result
	
	# 工具方法
	def update_status(self, status=None, progress=None):
		"""更新状态和进度"""
		if status is not None:
			self.status = status
			self.db_update()
			jingrow.db.commit()
			
		if progress is not None:
			self.progress = progress
			self.save()
			jingrow.db.commit()
	
	def is_valid_image_url(self, url):
		"""验证图片URL是否有效"""
		return url and url.strip() and url.lower() != "none"
	
	def get_reference_images(self):
		"""获取有效的参考图片列表"""
		images = []
		for i in range(1, 5):
			img_field = f"reference_image_{i}"
			if hasattr(self, img_field) and self.is_valid_image_url(getattr(self, img_field)):
				images.append(getattr(self, img_field))
		return images
	

# 工具函数
def get_context(pg):
	"""获取Ai Agent条件评估的上下文"""
	return {
		"pg": pg,
		"nowdate": nowdate,
		"utils": get_safe_globals().get("jingrow").get("utils"),
	}


def execute_agent_task(agent_name, agent_pagetype):
	"""后台执行智能体任务"""
	try:
		agent = jingrow.get_pg(agent_pagetype, agent_name)
		if not agent:
			return
		agent.execute_agent()
	except Exception as e:
		jingrow.log_error("Ai Agent任务执行异常", f"agent_name={agent_name}, agent_pagetype={agent_pagetype}, error={str(e)}")


@jingrow.whitelist()
def execute_agent_by_name(agent_name):
	"""执行指定的 Ai Agent（通过 RPC 调用）"""
	if not agent_name:
		return {"success": False, "error": "agent_name不能为空"}
	
	try:
		# 检查 agent 是否存在
		if not jingrow.db.exists('Ai Agent', agent_name):
			return {"success": False, "error": f"Ai Agent {agent_name} 不存在"}
		
		# 异步执行，不阻塞前端响应
		jingrow.enqueue(
			method=execute_agent_task,
			agent_name=agent_name,
			agent_pagetype='Ai Agent',
			queue='default',
			timeout=3600
		)
		
		return {
			"success": True,
			"message": "智能体执行任务已加入队列，请稍后查看执行状态"
		}
		
	except Exception as e:
		jingrow.log_error("Ai Agent执行失败", f"agent_name={agent_name}, error={str(e)}")
		return {"success": False, "error": str(e)}


def enqueue_ai_agent(pg, agent_name):
	"""后台执行Ai Agent任务"""
	try:
		# 通过 agent_name 获取完整的 Page 实例
		agent = jingrow.get_pg("Ai Agent", agent_name)
		if not agent:
			return
		
		agent_name = getattr(agent, "agent_name", None)
		initial_inputs = {
			"pagetype": pg.get("pagetype"),
			"name": pg.get("name"),
			"agent_name": agent_name
		}
		
		agent.execute_agent(initial_inputs=initial_inputs)
		
	except Exception as e:
		jingrow.log_error("Ai Agent执行失败", f"Ai Agent {agent_name} 执行失败: {str(e)}")


def execute_scheduled_agents(agent_name):
	"""定时任务统一入口，执行指定的Ai Agent"""
	try:
		if not agent_name:
			return
		
		agent = jingrow.get_pg("Ai Agent", agent_name)
		if not agent:
			return
		
		if not agent.enabled or agent.trigger_mode != "Scheduled Trigger":
			return
		
		agent.execute_agent()
	except Exception as e:
		jingrow.log_error("定时Ai Agent执行异常", f"agent_name={agent_name}, error={str(e)}")
