import socket
import subprocess
import tempfile
import json
import os
from typing import Dict
import time


class KatexError(Exception):
	pass

class KatexClient:
	def __init__(self):
		# Create temporary directory for socket
		self._temp_dir = tempfile.TemporaryDirectory(prefix='formatitko')
		self._socket_file = self._temp_dir.name + "/katex-socket"

		self._server_process = subprocess.Popen(["node", os.path.dirname(os.path.realpath(__file__)) + "/katex-server/index.mjs", self._socket_file])

		self._client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)

		# Wait for the node program to create the socket file
		while not os.path.exists(self._socket_file):
			time.sleep(0.01)

		# Wait for the node program to start accepting connections
		while True:
			try:
				self._client.connect(self._socket_file)
				time.sleep(0.01)
			except ConnectionRefusedError:
				continue
			break

	def render(self, tex: str, options: Dict={}):
		# Send formulas to translate
		self._client.sendall((json.dumps({"formulas":[{"tex":tex}], "options":options})+"\n").encode("utf-8"))

		# Receive response
		data = self._client.recv(4096)
		while data[-1] != 0x0a:
			data += self._client.recv(128)
		response = json.loads(data)

		if "error" in response:
			raise Exception(response["error"])
		if "error" in response["results"][0]:
			raise KatexError(response["results"][0]["error"])
		else:
			return response["results"][0]["html"]

	# Special commands implemented in the JS file for grouping defs together.
	def begingroup(self):
		self._client.sendall("begingroup\n".encode("utf-8"))

	def endgroup(self):
		self._client.sendall("endgroup\n".encode("utf-8"))

	def __enter__(self):
		return self

	def __exit__(self, type, value, tb):
		self._server_process.terminate()