
import os, sys
import json
import asyncio
import argparse
import fingerprint
from traceback import format_exc
from sqlor.dbpools import DBPools
from appPublic.jsonConfig import getConfig
from appPublic.log import debug, exception
from appPublic.worker import awaitify
from appPublic.dictObject import DictObject
from appPublic.timeUtils import curDateString, timestampstr
from appPublic.uniqueID import getID
from appPublic.folderUtils import listFile, _mkdir
from ahserver.filestorage import FileStorage
import av
from PIL import Image
from imageinfo import get_image_info, get_imagefaces

from pillow_heif import register_heif_opener
register_heif_opener()

def get_thumbnail(mloc, loctype, userid):
	fs = FileStorage()
	if loctype == '0': # webpath
		mloc1 = fs.realPath(mloc)
		mloc = mloc1
	# Open the original image
	with Image.open(mloc) as img:
		# Define the size of the thumbnail
		thumbnail_size = (256, 256)  # for example, 128x128 pixels
		# Create the thumbnail
		img.thumbnail(thumbnail_size)
		name = os.path.basename(mloc).split('.')[0] + '.jpg'
		fn = fs._name2path(name, userid=userid)
		webpath = fs.webpath(fn)
		_mkdir(os.path.dirname(fn))

		# Save the thumbnail to disk
		img.save(fn, 'JPEG')
		return webpath

def get_1st_keyframe(mloc, loctype,  userid):
	# Open the video file
	fs = FileStorage()
	if loctype == '0': # webpath
		mloc1 = fs.realPath(mloc)
		mloc = mloc1
	container = av.open(mloc)
	# Iterate through the video to find the first key frame
	kcnt = 0
	for frame in container.decode(video=0):
		if frame.key_frame:
			kcnt += 1
			if kcnt > 10:
				# Save the key frame as an image
				name = os.path.basename(mloc).split('.')[0] + '.jpg'
				fn = fs._name2path(name, userid=userid)
				webpath = fs.webpath(fn)
				_mkdir(os.path.dirname(fn))
				frame.to_image().save(fn)
				return webpath

async def save_image_info(sor, userid, imgid, path, ownerid):
	f = awaitify(get_image_info)
	lat, lon, ts = await f(path)
	finfo = await awaitify(get_imagefaces)(path, imgid, ownerid)
	ns1 = {
		"id":imgid
	}
	if ts is None:
		ns1["info_status"] = '0'
	else:
		ns1["info_status"] = "1"
		ns1["latitude"] = lat
		ns1["longitude"] = lon
		ns1["media_timestamp"] = ts
		ns1["face_cnt"] = len(finfo['faces'])
	await sor.C('imageinfo', ns1.copy())
	for f in finfo['faces']:
		ns2 = {
			'id': getID(),
			'imageid':imgid,
			'faceid':f['id']
		}
		await sor.C('imageface', ns2)

async def save_image(sor, path, webpath, ownerid):
	ns = DictObject()
	ns.id = getID()
	ns.name = os.path.basename(path).rsplit('.', 1)[0]
	ns.mtype = 'image'
	ns.mlocation = webpath
	ns.mloctype = '0'
	ns.logo_url = get_thumbnail(ns.mlocation, ns.mloctype,  ownerid)
	ns.create_date = curDateString()
	ns.ownerid = ownerid
	ns.del_flg = '0'
	refer_cnt = 1
	await sor.C('media', ns.copy())
	# await save_image_info(sor, ns.id, path, ownerid)
	return

async def save_audio(sor, path, webpath, ownerid):
	ns = DictObject()
	ns.id = getID()
	ns.name = os.path.basename(path).rsplit('.', 1)[0]
	ns.mtype = 'audio'
	ns.mlocation = webpath
	ns.mloctype = '0'
	ns.create_date = curDateString()
	ns.ownerid = ownerid
	ns.del_flg = '0'
	refer_cnt = 1
	await sor.C('media', ns.copy())

async def save_video(sor, path, webpath, ownerid):
	ns = DictObject()
	ns.id = getID()
	ns.name = os.path.basename(path).rsplit('.', 1)[0]
	ns.mtype = 'video'
	ns.mlocation = webpath
	ns.mloctype = '0'
	ns.logo_url = get_1st_keyframe(ns.mlocation, ns.mloctype,  ownerid)
	ns.create_date = curDateString()
	ns.ownerid = ownerid
	ns.del_flg = '0'
	refer_cnt = 1
	await sor.C('media', ns.copy())

ftyp_f = {
	"jpg":save_image,
	"jpeg":save_image,
	"heic":save_image,
	"gif":save_image,
	"png":save_image,
	"cr2":save_image,
	"mkv":save_video,
	"flv":save_video,
	"webm":save_video,
	"avi":save_video,
	"mov":save_video,
	"mp4":save_video,
	"wmv":save_video,
	"mp3":save_audio,
	"aiff":save_audio,
	"ape":save_audio,
	"wav":save_audio,
	"flac":save_audio,
	"aac":save_audio,
	"ogg":save_audio,
	"wma":save_audio,
	"alac":save_audio,
}
async def save_media(f, webpath, userid):
	fs = FileStorage()
	db = DBPools()
	async with db.sqlorContext('mediadb') as sor:
		if f is None:
			f = fs.realPath(webpath)
		recs = await sor.R('media', {'mlocation':webpath, 'mloctype':'0'})
		if len(recs) > 0:
			debug(f'{f} already exists')
			return

		sub = f.split('.')[-1].lower()
		func = ftyp_f.get(sub)
		if func:
			try:
				r = await func(sor, f, webpath, userid)
				debug(f'{f} loaded ...')
				return True
			except Exception as e:
				x = format_exc()
				debug(f'{f} is not loaded{e}\n{x}')
		else:
			debug(f'{f} is not loaded')
	return False

async def load_media(folder, username=None):
	folder = os.path.abspath(folder)
	debug(f'handle {folder} ...')
	fs = FileStorage()
	if not folder.startswith(fs.root):
		debug(f'{folder=} must start with {fs.root}')
		raise Exception(f'{folder=} must start with {fs.root}')
	userid = None
	if username:
		db = DBPools()
		async with db.sqlorContext('mediadb') as sor:
			users = await sor.R('users', {'username':username})
			if len(users) > 0:
				userid = users[0].id

	for f in listFile(folder, rescursive=True):
		webpath = fs.webpath(f)
		await save_media(f, webpath, userid)

if __name__ == '__main__':
	async def main():
		parser = argparse.ArgumentParser(prog=sys.argv[0])
		parser.add_argument('-w', '--workdir')
		parser.add_argument('-u', '--username')
		parser.add_argument('folders', nargs='*')
		args = parser.parse_args()
		workdir = args.workdir or os.getcwd()
		config = getConfig(workdir, {'workdir':workdir})
		DBPools(config.databases)
		for folder in args.folders:
			await load_media(folder, username=args.username or 'ymq')

	loop = asyncio.new_event_loop()
	loop.run_until_complete(main())

