import mautrix.appservice
import mautrix.types.event.message
import mautrix.types.event.type
import mautrix.api
import asyncio
import hashlib
import os
import asyncache, cachetools
import aiohttp
import re

# For Matrix, dynamic thumbnails and unauthorized media should be both enabled.

class GatewayMatrix:
	def __init__(self, plexus, settings, gateway_name, gateway_data):
		global PLEXUS
		PLEXUS = plexus
		
		self.appserv = mautrix.appservice.AppService(server = gateway_data["server"], domain = gateway_data["domain"], verify_ssl = False, id = gateway_data["id"], as_token = gateway_data["as_token"], hs_token = gateway_data["hs_token"], bot_localpart = gateway_data["bot_localpart"], bridge_name = gateway_data["bridge_name"], loop = asyncio.get_event_loop())
		
		self.settings = settings
		
		self.gateway_name = gateway_name
		
		self.bot_localpart = gateway_data["bot_localpart"]
		self.domain = gateway_data["domain"]
		
		self.room_name_to_room_id = {}
		self.room_id_to_room_name = {}
		
		self.xdm_id_to_matrix_id = {}
		self.matrix_id_to_xdm_id = {}
		
		self.msg_unique_id_to_event_id = {}
		self.event_id_to_msg_unique_id = {}
		self.msg_original_sender = {}
		
		self.my_puppets = {}
		
		self.active_users = {}
		
		self.ready = False
		
		def new_user(identification, nick_name):
			if identification[0] == "matrix" and identification[1] == self.gateway_name:
				# Ignore ours
				return
		PLEXUS.sub("new_user", new_user)
		
		def message(msg_unique_id, room_name, nick_name, avatar, gateway_type, gateway_name, unique_id, body, attachments, was_edit):
			if gateway_type == "matrix" and gateway_name == self.gateway_name:
				# Ignore ours
				return
			
			if len(attachments):
				body = str(body) + "\n\n" + "\n".join([str(a) for a in attachments])
			
			puppet = self.load_puppet((gateway_type, gateway_name, unique_id))
			
			# Not a problem if this gets called multiple times
			if room_name not in puppet["rooms"]:
				async def enter_room():
					await puppet["intent"].ensure_joined(room_id = self.room_name_to_room_id[room_name], ignore_cache = True)
					return True
				puppet["queue"].append(enter_room)
			
			if was_edit:
				async def q_msg():
					await puppet["intent"].send_message_event(room_id = self.room_name_to_room_id[room_name], event_type = mautrix.types.event.type.EventType.ROOM_MESSAGE, content = mautrix.types.event.message.TextMessageEventContent(msgtype = mautrix.types.event.message.MessageType.TEXT, format = mautrix.types.event.message.Format.HTML, formatted_body = body, relates_to = mautrix.types.event.message.RelatesTo(rel_type = mautrix.types.event.message.RelationType.REPLACE, event_id = self.msg_unique_id_to_event_id[msg_unique_id])))
					
					return True
				puppet["queue"].append(q_msg)
			else:
				async def q_msg():
					await puppet["intent"].set_displayname(displayname = nick_name, check_current = True)
					
					avatar_hash = hashlib.sha256(avatar).hexdigest() if (avatar is not None) else None
					if puppet["last_avatar_hash"] != avatar_hash:
						puppet["last_avatar_hash"] = avatar_hash
						await puppet["intent"].set_avatar_url(await puppet["intent"].upload_media(data = avatar))
					
					event_id = await puppet["intent"].send_message_event(room_id = self.room_name_to_room_id[room_name], event_type = mautrix.types.event.type.EventType.ROOM_MESSAGE, content = mautrix.types.event.message.TextMessageEventContent(msgtype = mautrix.types.event.message.MessageType.TEXT, format = mautrix.types.event.message.Format.HTML, formatted_body = body))
					
					self.msg_unique_id_to_event_id[msg_unique_id] = event_id
					self.event_id_to_msg_unique_id[event_id] = msg_unique_id
					self.msg_original_sender[msg_unique_id] = (gateway_type, gateway_name, unique_id)
					
					return True
				puppet["queue"].append(q_msg)
		PLEXUS.sub("message", message)
		
		def message_delete(msg_unique_id, room_name, gateway_type, gateway_name, unique_id):
			if gateway_type == "matrix" and gateway_name == self.gateway_name:
				# Ignore ours
				return
			
			puppet = self.load_puppet((gateway_type, gateway_name, unique_id))
			
			async def q_redact():
				await puppet["intent"].redact(room_id = self.room_name_to_room_id[room_name], event_id = self.msg_unique_id_to_event_id[msg_unique_id])
				
				return True
			puppet["queue"].append(q_redact)
		PLEXUS.sub("message_delete", message_delete)
		
		async def f():
			await self.appserv.start(host = "127.0.0.1", port = gateway_data["as_port"])
			
			for room_name, room_gateways in self.settings.ROOMS.items():
				for gateway in room_gateways:
					if "matrix" in gateway and gateway["matrix"] == self.gateway_name:
						room_id = await self.get_room_id_from_alias(gateway["room_alias"])
						
						self.room_name_to_room_id[room_name] = room_id
						self.room_id_to_room_name[room_id] = room_name
						
						await self.appserv.intent.ensure_joined(room_id = room_id, ignore_cache = True)
			
			self.appserv.matrix_event_handler(self.matrix_event_handler)
			
			self.ready = True
		
		asyncio.get_event_loop().create_task(f())
		asyncio.get_event_loop().create_task(self.step_puppets())
	
	async def matrix_event_handler(self, event):
		if not event.sender or not event.room_id or event.room_id not in self.room_id_to_room_name:
			return
		
		if re.match(re.compile(f"@{re.escape(self.bot_localpart)}[^:]*\\:{re.escape(self.domain)}"), event.sender):
			return
		
		if str(event.type) == "m.room.message":
			if str(event.content.msgtype) != "m.text":
				return
			
			was_edit = event.content._relates_to and str(event.content._relates_to.rel_type) == "m.replace"
			
			if was_edit:
				msg_unique_id = self.event_id_to_msg_unique_id[event.content._relates_to.event_id]
			else:
				msg_unique_id = os.urandom(16)
				self.msg_unique_id_to_event_id[msg_unique_id] = event.event_id
				self.event_id_to_msg_unique_id[event.event_id] = msg_unique_id
				self.msg_original_sender[msg_unique_id] = ("matrix", self.gateway_name, event.sender)
			
			room_name = self.room_id_to_room_name[event.room_id]
			nick_name = (await self.appserv.intent.get_displayname(event.sender)) or event.sender
			avatar = await self.get_native_user_avatar(event.sender)
			
			body = event.content.formatted_body if event.content.format else event.content.body
			attachments = []
			
			if ("matrix", self.gateway_name, event.sender) not in self.active_users:
				self.active_users[("matrix", self.gateway_name, event.sender)] = None
				PLEXUS.pub("new_user", ("matrix", self.gateway_name, event.sender), nick_name)
			
			PLEXUS.pub("message", msg_unique_id, room_name, nick_name, avatar, "matrix", self.gateway_name, event.sender, body, attachments, was_edit)
		elif str(event.type) == "m.room.redaction":
			if event.redacts not in self.event_id_to_msg_unique_id:
				return
			if self.event_id_to_msg_unique_id[event.redacts] not in self.msg_original_sender:
				return
			
			PLEXUS.pub("message_delete", self.event_id_to_msg_unique_id[event.redacts], self.room_id_to_room_name[event.room_id], *self.msg_original_sender[self.event_id_to_msg_unique_id[event.redacts]])
	
	def load_puppet(self, identification):
		if identification in self.my_puppets:
			return self.my_puppets[identification]
		
		puppet_mxid = f"@_uw_bridge_{hashlib.sha256(str(identification).encode('UTF-8')).hexdigest()[:24]}:underware.dev"
		
		puppet_data = {
			"mxid": puppet_mxid,
			"rooms": {},
			"queue": [],
			"last_avatar_hash": None,
		}
		
		async def get_intent():
			puppet_data["intent"] = self.appserv.intent.user(puppet_mxid)
			return True
		puppet_data["queue"].append(get_intent)
		
		self.my_puppets[identification] = puppet_data
		
		return puppet_data
	
	async def step_puppets(self):
		while not self.ready:
			await asyncio.sleep(2)
		
		while True:
			for xdm_id, puppet_data in self.my_puppets.items():
				if len(puppet_data["queue"]) > 0:
					if await puppet_data["queue"][0]():
						puppet_data["queue"].pop(0)
			await asyncio.sleep(0.2)
	
	# I couldn't find this in Mautrix itself, so wtf
	async def get_room_id_from_alias(self, alias):
		return (await self.appserv.intent.api.request(method = mautrix.api.Method.GET, path = mautrix.api.PathBuilder(f"_matrix/client/v3/directory/room/{alias}")))["room_id"]
	
	@asyncache.cached(cachetools.TTLCache(256, 1800))
	async def get_native_user_avatar(self, mxid):
		mxc_uri = await self.appserv.intent.get_avatar_url(mxid)
		
		if mxc_uri:
			# https://mementomori.social/@rolle/113732824034474418
			async with self.appserv.http_session.get(str(self.appserv.intent.api.get_download_url(mxc_uri, download_type = "thumbnail")) + "?width=128&height=128&method=scale&allow_redirect=true") as resp:
				return await resp.read()
		
		return None