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

import functools
import json
from datetime import datetime
import re

from rq.command import send_stop_job_command
from rq.exceptions import InvalidJobOperation, NoSuchJobError
from rq.job import Job
from rq.queue import Queue

import jingrow
from jingrow import _
from jingrow.model.page import Page
from jingrow.utils import (
	cint,
	compare,
	convert_utc_to_system_timezone,
	create_batch,
	make_filter_dict,
)
from jingrow.utils.background_jobs import get_queues, get_redis_conn

QUEUES = ["default", "long", "short"]
JOB_STATUSES = ["queued", "started", "failed", "finished", "deferred", "scheduled", "canceled"]

LOCAL_JOB_CACHE_KEY = "local_jobs"

def _coerce_redis_obj(obj: dict) -> jingrow._dict:
	modified = obj.get("modified") or obj.get("ended_at") or obj.get("started_at") or obj.get("creation")
	return jingrow._dict(
		name=obj.get("job_id") or obj.get("name"),
		job_id=obj.get("job_id"),
		queue=obj.get("queue"),
		job_name=obj.get("job_name") or obj.get("method"),
		status=(obj.get("status") or "queued").lower(),
		started_at=obj.get("started_at") or "",
		ended_at=obj.get("ended_at") or "",
		time_taken=obj.get("time_taken") or "",
		exc_info=obj.get("exc_info") or "",
		arguments=jingrow.as_json(obj.get("arguments") or obj.get("payload") or {}),
		timeout=obj.get("timeout"),
		creation=obj.get("creation") or modified,
		modified=modified,
		_comment_count=0,
		owner=obj.get("owner") or obj.get("user"),
		modified_by=obj.get("modified_by") or obj.get("user"),
	)


def _redis_set(job_id: str, data: dict) -> None:
	data = dict(data or {})
	data.setdefault("job_id", job_id)
	now = convert_utc_to_system_timezone(datetime.utcnow())
	data["modified"] = data.get("modified") or now
	
	# 设置默认的owner和creation字段
	data["owner"] = data.get("owner") or jingrow.session.user
	data["creation"] = data.get("creation") or data.get("modified") or now
	data["modified_by"] = data.get("modified_by") or jingrow.session.user
	
	if data.get("status") == "started" and not data.get("started_at"):
		data["started_at"] = now
	if data.get("status") in {"finished", "failed", "canceled"} and not data.get("ended_at"):
		data["ended_at"] = now
	cache = jingrow.cache()
	try:
		existing_raw = cache.hget(LOCAL_JOB_CACHE_KEY, job_id)
		if existing_raw:
			try:
				existing = json.loads(existing_raw)
			except Exception:
				existing = {}
			merged = {**existing, **data}
			data = merged
	except Exception:
		pass
	try:
		if isinstance(data.get("arguments"), (dict, list)):
			data["arguments"] = jingrow.as_json(data.get("arguments"))
	except Exception:
		pass
	cache.hset(LOCAL_JOB_CACHE_KEY, job_id, jingrow.as_json(data))


def _redis_get(job_id: str) -> dict:
	cache = jingrow.cache()
	raw = cache.hget(LOCAL_JOB_CACHE_KEY, job_id)
	if not raw:
		raise jingrow.DoesNotExistError
	try:
		return json.loads(raw)
	except Exception:
		return {}


def _redis_delete(job_id: str) -> None:
	cache = jingrow.cache()
	cache.hdel(LOCAL_JOB_CACHE_KEY, job_id)


def check_permissions(method):
	@functools.wraps(method)
	def wrapper(*args, **kwargs):
		jingrow.only_for("System Manager")
		instance = args[0]
		job = getattr(instance, "job", None)
		if job and (not for_current_site(job)):
			raise jingrow.PermissionError

		return method(*args, **kwargs)

	return wrapper


class LocalJob(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

		arguments: DF.Code | None
		ended_at: DF.Datetime | None
		exc_info: DF.Code | None
		job_id: DF.Data | None
		job_name: DF.Data | None
		queue: DF.Literal["default", "short", "long"]
		started_at: DF.Datetime | None
		status: DF.Literal["queued", "started", "finished", "failed", "deferred", "scheduled", "canceled"]
		time_taken: DF.Duration | None
		timeout: DF.Duration | None
	# end: auto-generated types

	def load_from_db(self):
		data = _redis_get(self.name)
		if not data:
			raise jingrow.DoesNotExistError
		super(Page, self).__init__(jingrow._dict(data))

	@property
	def job(self):
		return self._job_obj

	@staticmethod
	def get_list(args):
		# 优先使用新的参数名 limit_start 和 limit_page_length，向后兼容旧的参数名
		start = cint(args.get("limit_start") or args.get("start")) or 0
		page_length = cint(args.get("limit_page_length") or args.get("page_length")) or 20

		order_desc = "desc" in args.get("order_by", "")

		items = []
		try:
			cache = jingrow.cache()
			all_map = cache.hgetall(LOCAL_JOB_CACHE_KEY) or {}
			for _, val in (all_map.items() if isinstance(all_map, dict) else []):
				try:
					obj = json.loads(val)
					items.append(_coerce_redis_obj(obj))
				except Exception:
					continue
		except Exception:
			pass

		items = sorted(items, key=lambda j: j.modified, reverse=order_desc)
		return items[start : start + page_length]

	@staticmethod
	def get_matching_job_ids(args) -> list[str]:
		# 与 get_list 相同：统计站点缓存哈希内元素个数
		try:
			cache = jingrow.cache()
			all_map = cache.hgetall(LOCAL_JOB_CACHE_KEY) or {}
			return [None] * (len(all_map) if isinstance(all_map, dict) else 0)
		except Exception:
			return []

	@check_permissions
	def delete(self):
		# 从站点缓存删除，并清理 RQ 中对应的 job（如果存在且允许）
		try:
			_redis_delete(self.name or self.job_id)
		except Exception:
			pass
		try:
			# 非强制：如果能拿到 RQ job，就尝试删除
			job_id = self.job_id or self.name
			if job_id:
				job = Job.fetch(job_id, connection=get_redis_conn())
				job and job.delete()
		except Exception:
			pass

	@check_permissions
	def stop_job(self):
		try:
			send_stop_job_command(connection=get_redis_conn(), job_id=self.job_id)
		except InvalidJobOperation:
			jingrow.msgprint(_("Job is not running."), title=_("Invalid Operation"))

	@staticmethod
	def get_count(args) -> int:
		return len(LocalJob.get_matching_job_ids(args))

	# None of these methods apply to virtual job pagetype, overriden for sanity.
	@staticmethod
	def get_stats(args):
		return {}

	def db_insert(self, *args, **kwargs):
		# 创建：写入站点缓存哈希
		job_id = getattr(self, "job_id", None) or getattr(self, "name", None)
		if not job_id:
			jingrow.throw("Invalid Local Job", "job_id is required")
		_redis_set(job_id, self.as_dict())
		self.name = job_id

	def db_update(self, *args, **kwargs):
		# 更新：覆盖写入站点缓存哈希
		job_id = getattr(self, "job_id", None) or getattr(self, "name", None)
		if not job_id:
			jingrow.throw("Invalid Local Job", "job_id is required")
		_redis_set(job_id, self.as_dict())


@jingrow.whitelist()
def push_local_job(data: dict | None = None, **kwargs):
    """供本地 Jflow 推送任务状态使用。
    兼容两种调用方式：
    - POST { data: { job_id, status, ... } }
    - POST { job_id: ..., status: ... }
    """
    try:
        payload = data if isinstance(data, dict) else {}
        if kwargs:
            payload = {**payload, **kwargs}
        if not isinstance(payload, dict) or not payload:
            return {"success": False, "error": "invalid payload"}

        job_id = payload.get("job_id")
        if not job_id:
            return {"success": False, "error": "job_id required"}
        _redis_set(job_id, payload)
        return {"success": True}
    except Exception as e:
        jingrow.log_error("push_local_job_failed", str(e))
        return {"success": False, "error": str(e)}

def serialize_job(job: Job) -> jingrow._dict:
	modified = job.last_heartbeat or job.ended_at or job.started_at or job.created_at
	job_kwargs = job.kwargs.get("kwargs", {})
	job_name = job_kwargs.get("job_type") or str(job.kwargs.get("job_name"))
	if job_name == "jingrow.utils.background_jobs.run_pg_method":
		pagetype = job_kwargs.get("pagetype")
		pg_method = job_kwargs.get("pg_method")
		if pagetype and pg_method:
			job_name = f"{pagetype}.{pg_method}"

	# function objects have this repr: '<function functionname at 0xmemory_address >'
	# This regex just removes unnecessary things around it.
	if matches := re.match(r"<function (?P<func_name>.*) at 0x.*>", job_name):
		job_name = matches.group("func_name")

	return jingrow._dict(
		name=job.id,
		job_id=job.id,
		queue=job.origin.rsplit(":", 1)[1],
		job_name=job_name,
		status=job.get_status(),
		started_at=convert_utc_to_system_timezone(job.started_at) if job.started_at else "",
		ended_at=convert_utc_to_system_timezone(job.ended_at) if job.ended_at else "",
		time_taken=(job.ended_at - job.started_at).total_seconds() if job.ended_at else "",
		exc_info=job.exc_info,
		arguments=jingrow.as_json(job.kwargs),
		timeout=job.timeout,
		creation=convert_utc_to_system_timezone(job.created_at),
		modified=convert_utc_to_system_timezone(modified),
		_comment_count=0,
		owner=job.kwargs.get("user"),
		modified_by=job.kwargs.get("user"),
	)


def for_current_site(job: Job) -> bool:
	return job.kwargs.get("site") == jingrow.local.site


def filter_current_site_jobs(job_ids: list[str]) -> list[str]:
	site = jingrow.local.site

	return [j for j in job_ids if j.startswith(site)]


def _eval_filters(filter, values: list[str]) -> list[str]:
	if filter:
		operator, operand = filter
		return [val for val in values if compare(val, operator, operand)]
	return values


def fetch_job_ids(queue: Queue, status: str) -> list[str]:
	registry_map = {
		"queued": queue,  # self
		"started": queue.started_job_registry,
		"finished": queue.finished_job_registry,
		"failed": queue.failed_job_registry,
		"deferred": queue.deferred_job_registry,
		"scheduled": queue.scheduled_job_registry,
		"canceled": queue.canceled_job_registry,
	}

	registry = registry_map.get(status)
	if registry is not None:
		job_ids = registry.get_job_ids()
		return [j for j in job_ids if j]

	return []


@jingrow.whitelist()
def remove_failed_jobs():
	jingrow.only_for("System Manager")
	for queue in get_queues():
		fail_registry = queue.failed_job_registry
		failed_jobs = filter_current_site_jobs(fail_registry.get_job_ids())

		# Delete in batches to avoid loading too many things in memory
		conn = get_redis_conn()
		for job_ids in create_batch(failed_jobs, 100):
			for job in Job.fetch_many(job_ids=job_ids, connection=conn):
				job and fail_registry.remove(job, delete_job=True)


@jingrow.whitelist()
def stop_job(job_id):
	jingrow.get_pg("Local Job", job_id).stop_job()


@jingrow.whitelist()
def get_local_job_count():
	"""获取Local Job总数"""
	try:
		args = {
			'pagetype': 'Local Job',
			'filters': [],
			'fields': [],
			'distinct': False,
			'limit': None
		}
		count = LocalJob.get_count(args)
		return count
	except Exception as e:
		jingrow.log_error("get_local_job_count_failed", str(e))
		return 0


@jingrow.whitelist()
def get_local_job_list(limit_start=0, limit_page_length=20, order_by="modified desc"):
	"""获取Local Job列表 - 用于API调用"""
	jingrow.only_for("System Manager")
	try:
		args = jingrow._dict(
			limit_start=limit_start,
			limit_page_length=limit_page_length,
			order_by=order_by
		)
		return LocalJob.get_list(args)
	except Exception as e:
		jingrow.log_error("get_local_job_list_failed", str(e))
		return []
