#!/bin/env python

"""
this script assumes a layout like the following, in the pwd of your
employees
/Videos
	/shows
		/????p
			/ShowName
				/[Season ?? ]Episode ?? [ - TITLE].mkv
	/movies
		/????p
			/[A-Z0-9@]
				/MovieName.mkv
"""

from sys import argv, stderr, exit
from glob import iglob
from itertools import chain
import shlex
import re
import readline
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
from os import system
import socket

RESOLUTIONS=[(1280, 720), (1366, 768), (1920, 1080)]
#RESOLUTIONS=[(710, 480), (1136, 768)]

argv0 = 'reenc.py'

def usage():
	global argv0
	print(f'usage: {argv0} <host> <port> <srcdir>', file=stderr)
	exit(-1)

class JobError(Exception):
	pass
def send_cmd(reader, writer, cmd):
	lines = []

	writer.write(cmd.encode()+b'\n')
	writer.flush()
	while True:
		line = reader.readline()
		if line == '?\n':
			raise JobError(cmd)
		elif line == 'OKAY\n':
			return lines
		else:
			lines.append(line.strip())


@dataclass
class Episode:
	file: Path

	subs: Optional[str] = None
	title: Optional[str] = None
	season: int = 0
	number: int = 0

	def outname(self):
		name = f'Episode {self.number:02}'
		if self.title is not None:
			name = f'{name} - {self.title}'
		if self.season != 0:
			name = f'Season {self.season:02} {name}'
		
		return f'{name}.mkv'
	
	def outpath(self, res, show):
		wid, hei = res
		return Path('Videos/shows') / f'{hei}p' / show / self.outname()

	def generate_script(self, show, res):
		wid, hei = res
		outpath = self.outpath(res, show)
		inparg = shlex.quote(str(self.file))
		outarg = shlex.quote(str(outpath))
		dirarg = shlex.quote(str(outpath.parent))
		return '#!/bin/sh\n' f'mkdir -p {dirarg}\n' f'ffmpeg -loglevel error -i {inparg} -map 0 -c:a copy -c:s ass -c:v h264 -s {wid}x{hei} {outarg}\n'
		
	def slug(self):
		if self.season == 0:
			return f'E{self.number:02}'
		else:
			return f'S{self.season:02}E{self.number:02}'

def find_common_ends(files):
	files = map(lambda x: re.sub('\d', '#', x), files)
	needle = next(files)
	pfxlen = len(needle)
	sfxlen = len(needle)
	
	for file in files:
		while needle[:pfxlen] != file[:pfxlen]:
			pfxlen -= 1
		while needle[-sfxlen:] != file[-sfxlen:]:
			sfxlen -= 1
	
	return needle[:pfxlen], needle[-sfxlen:]

def find_sXXeXX(name):
	m = re.search('S([0-9]+)E([0-9]+)', name, re.I)
	if m is not None:
		season, episode = m.groups()
		return int(season), int(episode)
	
	season, episode = 0, 0
	
	if (m := re.search('S([0-9]+)', name, re.I)) is not None:
		season = int(m.group(1))
		name = name[:m.start()] + name[m.end():]
	elif (m := re.search('Season ([0-9]+)', name, re.I)) is not None:
		season = int(m.group(1))
		name = name[:m.start()] + name[m.end():]
		
	
	if (m := re.search('E([0-9]+)', name, re.I)) is not None:
		episode = int(m.group(1))
	elif (m := re.search('Episode ([0-9]+)', name, re.I)) is not None:
		episode = int(m.group(1))
	elif (m := re.search('([0-9]+)', name, re.I)) is not None:
		episode = int(m.group(1))
	
	return season, episode

def dump_metafile(path, name, files):
	with open(path, 'wt') as fp:
		print(f'showname: {name}',file=fp)
		print(file=fp)
		
		pfx, sfx = find_common_ends(map(lambda x: x.name, files))
		
		for file in files:
			season, episode = find_sXXeXX(file.name)
			title = file.name[len(pfx):-len(sfx)]

			print(f'file: {str(file)}', file=fp)
			print(f'season: {season}', file=fp)
			print(f'number: {episode}', file=fp)
			print(f'title: {title}', file=fp)
			print(file=fp)

def load_metafile(path):
	with open(path, 'rt') as fp:
		lines = map(str.strip, filter(lambda x: not x.startswith('#'), fp.readlines()))
		
		showname = None
		
		for line in lines:
			if line == '':
				break
			if ':' not in line:
				continue
			k,v = line.split(':', 1)
			
			if k.strip() == 'showname':
				showname = v.strip()


		episodes = []
		data = {}

		for line in lines:
			if line == '':
				if len(data) > 0:
					episodes.append(Episode(**data))
					data = {}
			if ':' not in line:
				continue
			k,v = line.split(':', 1)
			
			try:
				v = int(v.strip())
			except:
				v = v.strip()
				if v == '':
					continue
			
			data[k.strip()] = v

		if len(data) != 0:
			episodes.append(Episode(**data))

		return showname, episodes

def check_overlap(episodes):
	overmap = {}
	ret = False
	for episode in episodes:
		name = episode.outname()
		if name in overmap:
			overmap[name].append(episode)
			ret = True
		else:
			overmap[name] = [episode]
	
	for path, eps in overmap.items():
		if len(eps) > 1:
			print(f'multiple files would be saved at {path}')
			for ep in eps:
				print(f'\t{ep.file}')
	
	return ret

def do_connect(host, port):
	sok = socket.create_connection((host, port))
	return sok.makefile('r'), sok.makefile('wb')

def queue_job(reader, writer, show, res, episode):
	jobname = f'{show} {episode.slug()} @ {res}'
	script = episode.generate_script(show, res).encode()
	
	writer.write(f'NEW JOB {len(script)}\n'.encode())
	writer.write(script)
	writer.flush()
	jid = reader.readline().strip()
	if jid == '?\n':
		print(f'couldn\'t queue {jobname}')
		return
	else:
		jid = int(jid)
	reader.readline()

	send_cmd(reader, writer, f'META name IS {jobname}')
	send_cmd(reader, writer, 'MARK AS available')

	print(f'{jobname} queued as {jid}')


if __name__ == '__main__':
	if len(argv) > 0:
		argv0 = argv[0]
		argv = argv[1:]
	if len(argv) != 3:
		usage()
	host, port, dir = argv


	print(f'reencoding {dir}...')

	dir = Path(dir)
	metapath = Path(dir) / 'reenc.meta'
	
	if metapath.exists():
		showname, episodes = load_metafile(metapath)
		if check_overlap(episodes):
			print("exiting.")
			exit(-2)
		
		reader, writer = do_connect(host, port)
		
		for episode in episodes:
			for res in RESOLUTIONS:
				if episode.outpath(res, showname).exists():
					print(f'{episode.slug()} @ {res} already exists, skipping')
			
				queue_job(reader, writer, showname, res, episode)
		
		writer.write(b'QUIT\n')
		writer.flush()

	else:
		print(f'no metafile found, generating {metapath}...')

		mkglob = lambda x: iglob(f'**/*.{x}', root_dir=dir, recursive=True)
		globs = chain(mkglob('mkv'), mkglob('mp4'), mkglob('mov'), mkglob('avi'))
		files = [*map(lambda p: dir / p, globs)]
		files.sort()

		showname = input('enter the name of this show\n% ')
		
		dump_metafile(metapath, showname, files)
		system(f'xdg-open {shlex.quote(str(metapath))}')