from contextlib import nullcontext
from time import sleep

import kombu
import kombu.exceptions

from celery.result import AsyncResult, GroupResult
from greenlet import GreenletExit
from rich.panel import Panel
from rich.padding import Padding

from rich.progress import Progress as RichProgress, SpinnerColumn, TextColumn, TimeElapsedColumn
from secator.config import CONFIG
from secator.definitions import STATE_COLORS
from secator.output_types import Error
from secator.rich import console
from secator.utils import debug, traceback_as_string


class CeleryData(object):
	"""Utility to simplify tracking a Celery task and all of its subtasks."""

	def iter_results(
			result,
			ids_map={},
			description=True,
			refresh_interval=CONFIG.runners.poll_frequency,
			print_remote_info=True,
			print_remote_title='Results'
		):
		"""Generator to get results from Celery task.

		Args:
			result (Union[AsyncResult, GroupResult]): Celery result.
			description (bool): Whether to show task description.
			refresh_interval (int): Refresh interval.
			print_remote_info (bool): Whether to display live results.
			print_remote_title (str): Title for the progress panel.

		Yields:
			dict: Subtasks state and results.
		"""
		# Display live results if print_remote_info is set
		if print_remote_info:
			class PanelProgress(RichProgress):
				def get_renderables(self):
					yield Padding(Panel(
						self.make_tasks_table(self.tasks),
						title=print_remote_title,
						border_style='bold gold3',
						expand=False,
						highlight=True), pad=(2, 0, 0, 0))
			progress = PanelProgress(
				SpinnerColumn('dots'),
				TextColumn('{task.fields[descr]}  ') if description else '',
				TextColumn('[bold cyan]{task.fields[full_name]}[/]'),
				TextColumn('{task.fields[state]:<20}'),
				TimeElapsedColumn(),
				TextColumn('{task.fields[count]}'),
				TextColumn('{task.fields[progress]}%'),
				# TextColumn('\[[bold magenta]{task.fields[id]:<30}[/]]'),  # noqa: W605
				refresh_per_second=1,
				transient=False,
				console=console,
				# redirect_stderr=True,
				# redirect_stdout=False
			)
		else:
			progress = nullcontext()

		with progress:

			# Make initial progress
			if print_remote_info:
				progress_cache = CeleryData.init_progress(progress, ids_map)

			# Get live results and print progress
			for data in CeleryData.poll(result, ids_map, refresh_interval):
				yield from data['results']

				if print_remote_info:
					task_id = data['id']
					progress_id = progress_cache[task_id]
					CeleryData.update_progress(progress, progress_id, data)

			# Update all tasks to 100 %
			if print_remote_info:
				for progress_id in progress_cache.values():
					progress.update(progress_id, advance=100)

	@staticmethod
	def init_progress(progress, ids_map):
		cache = {}
		for task_id, data in ids_map.items():
			pdata = data.copy()
			state = data['state']
			pdata['state'] = f'[{STATE_COLORS[state]}]{state}[/]'
			id = progress.add_task('', advance=0, **pdata)
			cache[task_id] = id
		return cache

	@staticmethod
	def update_progress(progress, progress_id, data):
		"""Update rich progress with fresh data."""
		pdata = data.copy()
		state = data['state']
		pdata['state'] = f'[{STATE_COLORS[state]}]{state}[/]'
		pdata = {k: v for k, v in pdata.items() if v}
		progress.update(progress_id, **pdata)

	@staticmethod
	def poll(result, ids_map, refresh_interval):
		"""Poll Celery subtasks results in real-time. Fetch task metadata and partial results from each task that runs.

		Yields:
			dict: Subtasks state and results.
		"""
		while True:
			try:
				yield from CeleryData.get_all_data(result, ids_map)
				if result.ready():
					debug('result is ready', sub='celery.poll', id=result.id)
					yield from CeleryData.get_all_data(result, ids_map)
					break
			except (KeyboardInterrupt, GreenletExit):
				debug('encounted KeyboardInterrupt or GreenletExit', sub='celery.poll')
				raise
			except Exception as e:
				error = Error.from_exception(e)
				debug(repr(error), sub='celery.poll')
				pass
			finally:
				sleep(refresh_interval)

	@staticmethod
	def get_all_data(result, ids_map):
		"""Get Celery results from main result object, AND all subtasks results.

		Yields:
			dict: Subtasks state and results.
		"""
		task_ids = list(ids_map.keys())
		for task_id in task_ids:
			data = CeleryData.get_task_data(task_id, ids_map)
			if not data:
				continue
			debug(
				'POLL',
				sub='celery.poll',
				id=data['id'],
				obj={data['full_name']: data['state'], 'count': data['count']},
				verbose=True
			)
			yield data

		# Calculate and yield parent task progress
		# if not datas:
		# 	return
		# total = len(datas)
		# count_finished = sum([i['ready'] for i in datas if i])
		# percent = int(count_finished * 100 / total) if total > 0 else 0
		# parent_id = [c for c in ids_map.values() if c['full_name'] == datas[-1]]
		# data['progress'] = percent
		# yield data

	@staticmethod
	def get_task_data(task_id, ids_map):
		"""Get task info.

		Args:
			task_id (str): Celery task id.

		Returns:
			dict: Task info (id, name, state, results, chunk_info, count, error, ready).
		"""

		# Get task data
		data = ids_map.get(task_id, {})
		if not data:
			ids_map[task_id] = {}
		elif data.get('ready', False):
			return

		# if not data:
		# 	debug('task not in ids_map', sub='debug.celery', id=task_id)
		# 	return

		# Get remote result
		res = AsyncResult(task_id)
		if not res:
			debug('empty response', sub='celery.data', id=task_id)
			return

		# Set up task state
		data.update({
			'state': res.state,
			'ready': False,
			'results': []
		})

		# Get remote task data
		info = res.info

		# Depending on the task state, info will be either an Exception (FAILURE), a list (SUCCESS), or a dict (RUNNING).
		# - If it's an Exception, it's an unhandled error.
		# - If it's a list, it's the task results.
		# - If it's a dict, it's the custom user metadata.

		if isinstance(info, Exception):
			debug('unhandled exception', obj={'msg': str(info), 'tb': traceback_as_string(info)}, sub='celery.data', id=task_id)
			raise info

		elif isinstance(info, list):
			data['results'] = info
			errors = [e for e in info if e._type == 'error']
			status = 'FAILURE' if errors else 'SUCCESS'
			data['count'] = len([c for c in info if c._source.startswith(data['name'])])
			data['state'] = status

		elif isinstance(info, dict):
			data.update(info)

		# Set ready flag and progress
		ready = data['state'] in ['FAILURE', 'SUCCESS', 'REVOKED']
		data['ready'] = ready
		ids_map[task_id]['ready'] = data['ready']
		if data['ready']:
			data['progress'] = 100
		elif data['results']:
			progresses = [e for e in data['results'] if e._type == 'progress' and e._source == data['full_name']]
			if progresses:
				data['progress'] = progresses[-1].percent

		debug('data', obj=data, sub='celery.data', id=task_id, verbose=True)
		return data

	@staticmethod
	def get_task_ids(result, ids=[]):
		"""Get all Celery task ids recursively.

		Args:
			result (Union[AsyncResult, GroupResult]): Celery result object.
			ids (list): List of ids.
		"""
		if result is None:
			return

		try:
			if isinstance(result, GroupResult):
				CeleryData.get_task_ids(result.parent, ids=ids)

			elif isinstance(result, AsyncResult):
				if result.id not in ids:
					ids.append(result.id)

			if hasattr(result, 'children'):
				children = result.children
				if isinstance(children, list):
					for child in children:
						CeleryData.get_task_ids(child, ids=ids)

			# Browse parent
			if hasattr(result, 'parent') and result.parent:
				CeleryData.get_task_ids(result.parent, ids=ids)

		except kombu.exceptions.DecodeError:
			debug('kombu decode error', sub='celery.data')
			return
