#!/bin/env python
"""
use clip to interrogate image(s)
"""

import io
import base64
import sys
import os
import asyncio

import filetype
from PIL import Image

from util import log, Map
import sdapi as sdapi


def decode(encoding):
    if encoding.startswith("data:image/"):
        encoding = encoding.split(";")[1].split(",")[1]
    return Image.open(io.BytesIO(base64.b64decode(encoding)))


def encode(f):
    image = Image.open(f)
    exif = image.getexif()
    with io.BytesIO() as stream:
        image.save(stream, 'JPEG', exif = exif)
        values = stream.getvalue()
        encoded = base64.b64encode(values).decode()
        return encoded


async def interrogate(f):
    if not filetype.is_image(f):
        log.info({ 'interrogate skip': f })
        return
    json = Map({ 'image': encode(f) })
    log.info({ 'interrogate': f })
    # run clip
    json.model = 'clip'
    res = await sdapi.post('/sdapi/v1/interrogate', json)
    # res = sdapi.postsync('/sdapi/v1/interrogate', json)
    caption = ""
    style = ""
    if 'caption' in res:
        caption = res.caption
        log.info({ 'interrogate caption': caption })
        if ', by' in caption:
            style = caption.split(', by')[1].strip()
            log.info({ 'interrogate style': style })
    else:
        log.error({ 'interrogate clip error': res })
    # run booru
    json.model = 'deepdanbooru'
    res = await sdapi.post('/sdapi/v1/interrogate', json)
    # res = sdapi.postsync('/sdapi/v1/interrogate', json)
    keywords = {}
    if 'caption' in res:
        for term in res.caption.split(', '):
            term = term.replace('(', '').replace(')', '').split(':')
            keywords[term[0]] = term[1]
        keywords = dict(sorted(keywords.items(), key=lambda x:x[1], reverse=True))
        log.info({ 'interrogate keywords': keywords })
    else:
        log.error({ 'interrogate booru error': res })
    return caption, keywords, style


async def main():
    sys.argv.pop(0)
    await sdapi.session()
    if len(sys.argv) == 0:
        log.error({ 'interrogate': 'no files specified' })
    for arg in sys.argv:
        if os.path.exists(arg):
            if os.path.isfile(arg):
                await interrogate(arg)
            elif os.path.isdir(arg):
                for root, _dirs, files in os.walk(arg):
                    for f in files:
                        await interrogate(os.path.join(root, f))
            else:
                log.error({ 'interrogate unknown file type': arg })
        else:
            log.error({ 'interrogate file missing': arg })
    await sdapi.close()


if __name__ == "__main__":
    asyncio.run(main())