#!/bin/env python3

"""
this software distributed under a "fuck around, find out" license. the
author gives nobody in the world any explicit legal right to use this
code for any purpose whatsoever. the author reserves the right to send
holy legal fury down from the heavens should you invoke her ire.

the author does not have the money to do so, in any juris diction
worldwide. the author is also prety chill, she thinks, so, you know, do
what you want or whatever. the text of this license is intentionally
god damn stupid specifically to scare away corporate lawyers and
attract internet queers

if you do choose to redistribute this software (to "fuck around", as it
were), you should change this disclaimer some and make it funnier

the mitochondrea is the powerhouse of the cell

"""

from sys import argv, stdin, stdout, stderr, exit, exc_info
import asyncio
import pickle
import gzip
import logging
from pathlib import Path
from dataclasses import dataclass, field, asdict
from datetime import datetime
from typing import Dict, Optional, List


MARKINGS = ('new', 'available', 'taken', 'done', 'failed')
argv0 = 'employer.py'


def usage():
	global argv0
	print(f'usage: python3 {argv0} <db path> [host] <port>', file=stderr)
	exit(-1)


@dataclass
class Job:
	jid: int
	script: bytes
	meta: Dict[str,str] = field(default_factory=dict)
	marking: str = 'new'

	created: datetime = field(default_factory=datetime.now)
	taken: Optional[datetime] = None
	completed: Optional[datetime] = None
	updated: Optional[datetime] = None
	
	def update(self):
		self.updated = datetime.now()
	
	def __getitem__(self, k):
		return self.meta.get(k, '')
	def __setitem__(self, k, v):
		if type(k) is not str or type(v) is not str:
			raise ValueError('key and value must be str')

		if v == '':
			del self.meta[k]
		else:
			self.meta[k] = v
		self.update()
	def __delitem__(self, k):
		del self.meta[k]
	
	def take(self):
		if self.marking != 'available':
			raise ValueError("can't take a job that's been taken")
		self.marking = 'taken'
		self.taken = datetime.now()
		
		self.update()

	def mark(self, marking):
		global MARKINGS
		if marking not in MARKINGS:
			raise ValueError(f"{marking}: not a marking")
		
		if self.marking == marking:
			return
		
		self.marking = marking
		if marking in ('new', 'available'):
			self.taken = None
			self.completed = None
		elif marking == 'taken':
			self.taken = datetime.now()
			self.completed = None
		else:
			self.completed = datetime.now()
		
		self.update()
	
	def infostr(self) -> str:
		s = f'SCRIPT IS {len(self.script)} BYTES\nMARKED AS {self.marking}\nCREATED AT {self.created.isoformat()}\n'
		if self.taken is not None:
			s += f'TAKEN AT {self.taken.isoformat()}\n'
		if self.completed is not None:
			s += f'COMPLETED AT {self.completed.isoformat()}\n'
		if self.updated is not None:
			s += f'UPDATED AT {self.updated.isoformat()}\n'
		return s
	def linestr(self) -> str:
		return f'{self.jid}\t{self.created.isoformat()}\t{len(self.script)}\t{self.marking}\n'
	def metastr(self) -> str:
		return ''.join(f'{k}: {v}\n' for (k, v) in self.meta.items())
	def script(self) -> bytes:
		return f'{len(self.script)}\n'.encode() + self.script

class Jobs:
	def __init__(self, filename):
		self.filename = filename
		self.lock = asyncio.Lock()

		try:
			fp = open(filename+'.lock', 'x')
			fp.write('')
		except FileExistsError:
			raise Exception(f"{filename} is already locked. kill the other instance or delete {filename}.lock.")

		if Path(filename).exists():
			data = pickle.load(gzip.open(filename, 'rb'))
			self.nextid = data['nextid']
			self.jobs = data['jobs']
		else:
			self.nextid = 0
			self.jobs = {}
			self.save()
	def __del__(self):
		Path(self.filename+'.lock').unlink()
	
	def __getitem__(self, k):
		global MARKINGS
		if type(k) is int:
			return self.jobs[k]
		elif k in MARKINGS:
			return filter(lambda j: j.marking == k, self.jobs.values())
		else:
			raise KeyError(k)
	def __delitem__(self, k):
		if type(k) is Job:
			del self.jobs[k.jid]
		else:
			del self.jobs[k]

	async def __aenter__(self):
		await self.lock.acquire()
	async def __aexit__(self, _exc_ty, _exc_val, _tb):
		self.save()
		self.lock.release()
	
	def save(self):
		data = {'nextid': self.nextid, 'jobs': self.jobs}
		with gzip.open(self.filename, 'wb') as fp:
			pickle.dump(data, fp)
	
	def new(self, script):
		job = Job(self.nextid, script)
		self.jobs[self.nextid] = job
		self.nextid += 1
		return job


def find_val_in_line(line) -> str:
	line = line[5:]
	return line[line.lower().index(b' is')+3:].strip().decode()

async def handle_line(topic, tokens, line, reader, writer) -> int:
	global jobs
	match tokens:
		case ['new', 'job', int(length)]:
			data = await reader.readexactly(length)
			job = jobs.new(data)
			writer.write(f'{job.jid}\n'.encode())
			return job.jid

		case ['list', str(marking), 'jobs']:
			writer.write(''.join(map(lambda x: x.linestr(), jobs[marking])).encode())
		case ['list', 'jobs']:
			writer.write(''.join(map(lambda x: x.linestr(), jobs.jobs.values())).encode())
		
		case ['info', 'for', 'job', int(topic), 'is']:
			writer.write(jobs[topic].infostr().encode())
		case ['info', 'is']:
			writer.write(jobs[topic].infostr().encode())

		case ['script', 'for', 'job', int(topic), 'is']:
			writer.write(f'{len(jobs[topic].script)}\n'.encode() + jobs[topic].script)
		case ['script', 'is']:
			writer.write(f'{len(jobs[topic].script)}\n'.encode() + jobs[topic].script)

		case ['script', 'for', 'job', int(topic), 'is', int(length)]:
			jobs[topic].script = await reader.readexactly(length)
		case ['script', 'is', int(length)]:
			jobs[topic].script = await reader.readexactly(length)

		case ['meta', 'for', 'job', int(topic), 'is']:
			writer.write(jobs[topic].metastr().encode())
		case ['meta', 'is']:
			writer.write(jobs[topic].metastr().encode())

		case ['meta', key, 'for', 'job', int(topic), 'is']:
			writer.write(f'{jobs[topic][key]}\n'.encode())
		case ['meta', key, 'is']:
			writer.write(f'{jobs[topic][key]}\n'.encode())
		
		case ['meta', key, 'for', 'job', int(topic), 'is', *_]:
			jobs[topic][key] = find_val_in_line(line)
		case ['meta', key, 'is', *_]:
			jobs[topic][key] = find_val_in_line(line)

		case ['meta', key, 'for', 'job', int(topic), 'clear']:
			del jobs[topic][key]
		case ['meta', key, 'clear']:
			del jobs[topic][key]
		
		case ['take', 'job', int(topic)]:
			jobs[topic].take()
		case ['take']:
			jobs[topic].take()
		
		case ['take', 'job']:
			try:
				avail = next(jobs['available'])
				avail.take()
				topic = avail.jid
				writer.write(f'{topic}\n'.encode())
			except StopIteration:
				raise ValueError('no new jobs')
		
		case ['mark', 'job', int(topic), 'as', str(marking)]:
			jobs[topic].mark(marking)
		case ['mark', 'as', str(marking)]:
			jobs[topic].mark(marking)
		
		case ['cancel', 'job', int(topic)]:
			del jobs[topic]
		case ['cancel']:
			del jobs[topic]

		case ['job', int(topic), 'exists']:
			jobs[topic]
		case ['exists']:
			jobs[topic]
		
		case ['quit']:
			writer.close()
		
		case _:
			raise ValueError(f'Unknown command {line!r}')
	
	return topic


def tokenize(line) -> List[bytes]:
	tokens = []
	it = filter(lambda x: x != '', line.decode().split(' '))
	for tok in it:
		try:
			tokens.append(int(tok))
		except ValueError:
			tokens.append(tok.lower())
	
	return tokens
async def accept(reader, writer):
	global jobs
	l = logging.getLogger('employer.sock')
	
	peer = writer.get_extra_info('peername')
	if peer is not None:
		peer = f'{peer[0]}:{peer[1]}'
	else:
		peer = 'unknown'
	
	l.info(f'connection from {peer}')
	
	topic = None
	
	try:
		while (line := await reader.readline()) != b'':
			line = line.strip(b'\n')
			if line == b'':
				continue

			tokens = tokenize(line)
			l.debug(f'{peer}: {line}, parsed as {tokens}')
			try:
				async with jobs:
					topic = await handle_line(topic, tokens, line, reader, writer)
				writer.write(b'OKAY\n')
			except (ValueError, KeyError):
				writer.write(b'?\n')
			except Exception:
				l.debug('', exc_info=exc_info())
				writer.write(b'?\n')
			
			await writer.drain()
	except ConnectionResetError:
		pass
	
	l.info(f'disconnect from {peer}')

async def main():
	global jobs, argv, argv0

	if len(argv) > 0:
		argv0 = argv[0]
		argv = argv[1:]
	if len(argv) not in (2,3):
		usage()

	file = argv[0]
	host = '' if len(argv) == 2 else argv[1]
	port = argv[-1]

	jobs = Jobs(file)

	l = logging.getLogger('employer')
	l.info(f'boot: using db {file}, serving {host}:{port}')

	server = await asyncio.start_server(accept, host, port)
	await server.serve_forever()


if __name__ == '__main__':
	logging.basicConfig(stream=stderr, level=logging.DEBUG, format='[%(asctime)s %(levelname)s %(name)s] %(msg)s')
	asyncio.run(main())