# Agenten Plattform
#
# (c) 2024 Magnus Bender
# 	Institute of Humanities-Centered Artificial Intelligence (CHAI)
# 	Universitaet Hamburg
# 	https://www.chai.uni-hamburg.de/~bender
#  
# source code released under the terms of GNU Public License Version 3
# https://www.gnu.org/licenses/gpl-3.0.txt

import random

from typing import Callable, List
from ums.agent import ExtractAudioAgent, ExtractImageAgent, ExtractTextAgent, SolveAgent, GatekeeperAgent

from ums.utils import AgentMessage, Riddle, RiddleData, RiddleSolution, RiddleStatus, ExtractedData

"""
	Examples for simple agents.

	Each agent is represented by its own class. The handling of tasks is done by `handle()` in each agent.

	Finally `AGENT_CLASSES` contains the classes of the agents in a list. Via environmental variables this list is specified to the ums.agent system. 
"""

class MyExtractAudioAgent(ExtractAudioAgent):

	def handle(self, data: RiddleData) -> RiddleData:
		print("Audio Process:", data.file_plain)
		return data

class MyExtractImageAgent(ExtractImageAgent):

	def handle(self, data: RiddleData) -> RiddleData:
		print("Image Process:", data.file_plain)

		extracted = ExtractedData(other={"info":"just a test"})
		data.file_extracted = self.store_extracted(data, extracted)
		return data
	
class MyExtractTextAgent(ExtractTextAgent):

	def before_response(self, response: AgentMessage, send_it: Callable[[], None]) -> bool:
		print("The response will be:", response)
		return True

	def handle(self, data: RiddleData) -> RiddleData:
		print("Text Process:", data.file_plain)
		return data


class MySolveAgent(SolveAgent):

	def handle(self, riddle: Riddle, data: List[RiddleData]) -> RiddleSolution:

		for d in data:
			print(self.get_extracted(d))
		
		if self.message().id == "test":
			status = RiddleStatus()
			status.extract.required = False
			self.sub_riddle(riddle=Riddle(context="Haha", question="Blubber"), status=status)

		return RiddleSolution(solution="Huii", explanation=f"Blubb, {random.random()}")
	

class MyGatekeeperAgent(GatekeeperAgent):

	def handle(self, solution: List[RiddleSolution], riddle: Riddle) -> RiddleSolution:
		solution[0].accepted = True 
		solution[0].review = "Ok"
		
		return solution

AGENT_CLASSES = [
	MyExtractAudioAgent, MyExtractImageAgent, MyExtractTextAgent,
	MySolveAgent,
	MyGatekeeperAgent
]