import socket
import subprocess
import tempfile
import json
import os

class KatexError(Exception):
	pass

class NPMNotFoundError(Exception):
	pass

class KatexServerError(Exception):
	pass


class KatexClient:
	_client: socket.socket
	_server_process: subprocess.Popen[bytes]
	_socket_file: str
	_temp_dir: tempfile.TemporaryDirectory[str]
	_connected: bool

	def __init__(self, socket: str=None, connect: bool=True):
		if socket is not None:
			self._socket_file = socket
		else:
			self.open_socket()
		if connect:
			self.connect()
			self._client.sendall("init\n".encode("utf-8")) # Reinitialize KaTeX Server in case it was reused.
			self._connected = True
		else:
			self._connected = False

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

		srcdir = os.path.dirname(os.path.realpath(__file__))

		# Test if `node_modules` directory exists and if not, run `npm install`
		if not os.path.isdir(srcdir + "/katex-server/node_modules"):
			print("Installing node dependencies for the first time...")
			try:
				subprocess.run(["npm", "install"], cwd=srcdir+"/katex-server", check=True)
			except subprocess.CalledProcessError as e:
				if e.returncode == 127:
					raise NPMNotFoundError("npm not found. Node.js is required to use KaTeX.")
				else:
					raise e

		self._server_process = subprocess.Popen(["node", srcdir + "/katex-server/index.mjs", self._socket_file], stdout=subprocess.PIPE)

		ok = self._server_process.stdout.readline()
		if ok != b"OK\n":
			raise KatexServerError("Failed to connect to katex-server")

	def connect(self):
		self._client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
		self._client.connect(self._socket_file)

	def get_socket(self):
		return self._socket_file

	def render(self, tex: str, options: dict={}):
		if not self._connected:
			raise KatexServerError("KatexClient not connected to Katex server. It should be initialized with connect=True.")
		# 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 KatexServerError(response["error"])
		if "error" in response["results"][0]:
			raise KatexError(response["results"][0]["error"] + " in $" + tex + "$")
		else:
			return response["results"][0]["html"]

	# Special commands implemented in the JS file for grouping defs together.
	def begingroup(self):
		if not self._connected:
			raise KatexServerError("KatexClient not connected to Katex server. It should be initialized with connect=True.")
		self._client.sendall("begingroup\n".encode("utf-8"))

	def endgroup(self):
		if not self._connected:
			raise KatexServerError("KatexClient not connected to Katex server. It should be initialized with connect=True.")
		self._client.sendall("endgroup\n".encode("utf-8"))

	def __enter__(self):
		return self

	def __exit__(self, type, value, tb):
		if hasattr(self, "_server_process") and self._server_process is not None:
			self._server_process.terminate()