import random

from abc import abstractmethod, ABC
from enum import Enum
from typing import List, Callable

from ums.utils import (
	RiddleInformation, AgentMessage, RiddleDataType, RiddleData, Riddle,
	RiddleStatus, RiddleSolution,
	logger
)

class AgentCapability(Enum):
	"""
		The three different capabilities an agent can have.
	"""

	EXTRACT="extract"
	SOLVE="solve"
	GATEKEEPER="gatekeeper"


class BasicAgent(ABC):
	"""
		A basic agent, each agent will be a subclass of this class.
	"""

	@staticmethod
	@abstractmethod
	def agent_capability() -> AgentCapability:
		"""
			Represents the capabilities of this agent, for messages/ tasks of this capability, the `handle` method will be called.
		"""
		pass

	def __init__(self, message:AgentMessage, send_message:Callable[[AgentMessage], bool]):
		self._send_message = send_message
		self._sub_cnt = 0

		self._message = message
		self._response = message.model_copy(deep=True)
		
		self._do_response = False

		self._process()
		self._respond()

	@abstractmethod
	def _process(self):
		pass

	def _respond(self):
		send_it = lambda: self._send_message(self._response)
		if self.before_response(self._response, send_it) and self._do_response:
			send_it()
			logger.debug(f"Response sent {self._response.id}")
		else:
			logger.debug(f"Stopped response {self._response.id}")
		
	def before_response(self, response:AgentMessage, send_it:Callable[[], None]) -> bool:
		"""
			This method is called before the response is sent.
			If the method returns `False` no response will be sent. 
			Thus, by overwriting this method, a response can be prevented.

			The response to be sent is in `response` and `send_it` is a callable, which sends the response to the management if it gets called. 
			(Hence, one may stop sending the response and later call `send_it()` to send the response.)
		"""
		return True
	
	def message(self) -> AgentMessage:
		"""
			Get the message this agent object is working on. 
		"""
		return self._message;

	def sub_riddle(self,
			riddle:Riddle, data:List[RiddleData]=[], status:RiddleStatus=None
		) -> AgentMessage|bool:
			"""
				Create a new sub-riddle for solving details of the current riddle.
				For the sub-riddle, give a `riddle` and optionally, a selection of `data` items (default none) and a `status` (default `ums.utils.types.RiddleStatus()`).
				By changing a status, different steps can be (de-)selected.

				Return the message of the sub-riddle or `false` on error.
			"""

			if status is None:
				status = RiddleStatus()
			
			# new sub riddle id
			self._sub_cnt += 1
			new_id = "{}-sub-{}.{}".format(self._message.id, self._sub_cnt, int(random.random()*100))

			self._message.sub_ids.append(new_id)
			self._response.sub_ids.append(new_id)

			# create the riddle's message
			sub_msg = AgentMessage(
				id=new_id,
				riddle=riddle,
				data=data,
				status=status
			)
			logger.debug(f"Created sub-riddle {sub_msg.id}")

			# send it
			if self._send_message(sub_msg):
				return sub_msg
			else:
				return False

	@abstractmethod
	def handle(self, *args:RiddleInformation) -> RiddleInformation:
		"""
			Handle a single task of the agent, the arguments and return value depends on the actual task (see subclass)!

			**This is the method to implement!**

			The full message is available via `message()`, a sub riddle can be created with `sub_riddle()`.
		"""
		pass

class ExtractAgent(BasicAgent):
	"""
		An extraction agent.
	"""

	def agent_capability() -> AgentCapability:
		return AgentCapability.EXTRACT
	
	@staticmethod
	@abstractmethod
	def extract_type() -> RiddleDataType:
		"""
			Represents the data this agent can process.
		"""
		pass

	def _process(self):
		for i, data in enumerate(self._response.data):
			if data.type == self.__class__.extract_type():
				logger.debug(f"Start extraction '{data.file_plain}'")
				result = self.handle(data)
				logger.debug(f"End extraction '{data.file_plain}' ('{result.file_extracted}')")

				if result.file_extracted is None:
					logger.info(f"Riddle {self._response.id}: 'file_extracted' for data '{data.file_plain}' still empty after handling")

				self._response.data[i] = result
				self._do_response = True

		self._response.status.extract.finished = True

	@abstractmethod
	def handle(self, data:RiddleData) -> RiddleData:
		"""
			Process the item `data`, create extraction file and return `data` with populated `data.file_extracted`.
		"""

class ExtractTextAgent(ExtractAgent):
	"""
		An extraction agent for text, create a subclass for your agent.
	"""
	
	def extract_type() -> RiddleDataType:
		return RiddleDataType.TEXT

class ExtractAudioAgent(ExtractAgent):
	"""
		An extraction agent for audio, create a subclass for your agent.
	"""
	
	def extract_type() -> RiddleDataType:
		return RiddleDataType.AUDIO

class ExtractImageAgent(ExtractAgent):
	"""
		An extraction agent for images, create a subclass for your agent.
	"""
	
	def extract_type() -> RiddleDataType:
		return RiddleDataType.IMAGE


class SolveAgent(BasicAgent):
	"""
		A solve agent, create a subclass for your agent.
	"""

	def agent_capability() -> AgentCapability:
		return AgentCapability.SOLVE
	
	def _process(self):
		logger.debug(f"Start solve: {self._response.id}")
		solution = self.handle(self._response.riddle, self._response.data)
		logger.debug(f"End solve: {self._response.id} ({solution.solution}, {solution.explanation})")
		
		if len(solution.solution) == 0 or len(solution.explanation) == 0:
			logger.info(f"Riddle {self._response.id}: Empty solution/ explanation after handling")

		self._response.solution = solution
		self._response.status.solve.finished = True

		self._do_response = True
	
	@abstractmethod
	def handle(self, riddle:Riddle, data:RiddleData) -> RiddleSolution:
		"""
			Solve the `riddle` using `data` and return a solution.
		"""
	
class GatekeeperAgent(BasicAgent):
	"""
		A gatekeeper agent, create a subclass for your agent.
	"""

	def agent_capability() -> AgentCapability:
		return AgentCapability.GATEKEEPER
	
	def _process(self):
		if self._response.solution is None:
			self._response.solution = RiddleSolution(solution="", explanation="")

		logger.debug(f"Start validate: {self._response.id}")
		solution = self.handle(self._response.solution, self._response.riddle)
		logger.debug(f"End validate: {self._response.id} ({solution.review}, {solution.accepted})")
		
		if solution.review is None or len(solution.review) == 0:
			logger.info(f"Riddle {self._response.id}: Empty review after handling")

		self._response.solution = solution
		self._response.status.validate.finished = True
		self._response.status.solved = solution.accepted

		self._do_response = True
	
	@abstractmethod
	def handle(self, solution:RiddleSolution, riddle:Riddle) -> RiddleSolution:
		"""
			Check the `solution` of `riddle` and return solution with populated `solution.accepted` and `solution.review`.
		"""