import os, logging
DEBUG = os.environ.get('DEBUG')
logging.basicConfig(
        #filename='protovax.log', encoding='utf-8',
        format='[%(asctime)s] %(levelname)s %(module)s/%(funcName)s - %(message)s',
        level=logging.DEBUG if DEBUG else logging.INFO)

import asyncio
import requests
import random
import json
import base64
import time
import re
from PIL import Image
from datetime import datetime, timezone, timedelta
from io import BytesIO
from telethon import TelegramClient, events

import secrets

PROTOSPACE = -1001298858490
JASON = 172033414
TANNER = 79316791
TANNERBOT_MICROSCOPY = -1001886549266

scope_subject = ''

client = TelegramClient('data/dreamer.session', secrets.TG_API_ID, secrets.TG_API_HASH)

def snap_size(size):
    VALID_SIZES = [512, 576, 640, 704, 768, 832, 896, 960, 1024]
    for s in VALID_SIZES:
        if s >= size:
            return s

def dreamer(prompt, steps=30, seed=1234, height=512, width=512):
    url = secrets.DREAM_API

    #sd_settings_json = '{"fn_index":12,"data":["","","None","None",30,"Euler a",false,false,1,1,7,-1,-1,0,0,0,false,512,512,false,false,0.7,"None",false,"Seed","","Steps","",true,false,false,null,"",null,"",""],"session_hash":"qpydye2zzsh"}'
    sd_settings_json = '{"fn_index":12,"data":["","","None","None",30,"Euler a",false,false,1,1,7,-1,-1,0,0,0,false,512,512,false,false,0.7,"None",false,false,null,"","Seed","","Nothing","",true,false,null,"",""],"session_hash":"qpydye2zzsh"}'

    sd_settings = json.loads(sd_settings_json)
    sd_settings['data'][0] = prompt
    sd_settings['data'][4] = steps
    sd_settings['data'][11] = seed
    sd_settings['data'][17] = snap_size(height)
    sd_settings['data'][18] = snap_size(width)

    if width > 512 or height > 512:
        # highres fix
        sd_settings['data'][19] = True

    logging.info('Dreaming of: %s', prompt)

    try:
        res = requests.post(url, json=sd_settings, timeout=40)
        res.raise_for_status()
        res = res.json()
        return res
    except BaseException as e:
        logging.error('Problem dreaming {}: {} - {}'.format(url, e.__class__.__name__, str(e)))
        return None

async def message_handler(event):
    global scope_subject

    text = event.raw_text
    chat = event.chat_id
    
    sender = await event.get_sender()
    if not event.sender:
        return

    if event.message.fwd_from:
        return

    name = getattr(sender, 'first_name', None)
    if not name:
        return

    command = ''
    data = ''
    reply = ''
    reply_file = None

    if text.startswith('! '):
        text = text[2:]
    elif text.startswith('!'):
        text = text[1:]
    else:
        return

    if ' ' in text:
        command = text.split(' ', 1)[0]
        data = text.split(' ', 1)[1]
    else:
        command = text

    command = command.lower()

    #!setscope
    if command == 'setscope' and data:
        if not reply and sender.id != TANNER:
            reply = '> do !scope to look under the microscope'
            await event.reply(reply)

        if not reply:
            scope_subject = data

            now = datetime.now(tz=timezone.utc) - timedelta(seconds=1)
            reply = '> ' + scope_subject
            reply_file = 'frames/' + now.strftime('image%S.jpg')
            reply_msg = await event.reply(reply, file=reply_file, link_preview=False)

            await reply_msg.forward_to(TANNERBOT_MICROSCOPY)

    #!scope
    if command == 'scope':
        logging.info('Chat: {:<14} | ID: {:<6} | User: {} {} | Command: {} | Data: {}'.format(chat, event.id, name, sender.id, command, data or 'None'))

        if not reply:

            now = datetime.now(tz=timezone.utc) - timedelta(seconds=1)
            reply = '> ' + scope_subject
            reply_file = 'frames/' + now.strftime('image%S.jpg')
            reply_msg = await event.reply(reply, file=reply_file, link_preview=False)


    #!dream
    if command == 'dream':
        logging.info('Chat: {:<14} | ID: {:<6} | User: {} {} | Command: {} | Data: {}'.format(chat, event.id, name, sender.id, command, data or 'None'))

        if not reply and chat == PROTOSPACE:
            reply = '> please join {} to dream'.format(secrets.PROTODREAM_URL)

        #if not reply and sender.id == JASON:
        #    reply = '> no.'

        to_send = ''
        if not reply and data:
            to_send = data

        if not reply and not to_send and event.message.is_reply:
            reply_msg = await event.get_reply_message()
            to_send = reply_msg.raw_text

        if not reply and '> ' in to_send:
            reply = '> sounds like a bad !dream'

        if not reply and len(to_send) > 500:
            reply = '> impressive. very nice. let\'s see Paul Allen\'s prompt.'

        if not reply and not to_send:
            reply = '> usage: !dream [prompt] [options]'
            reply += '\n> or reply to a someone with !dream'
            reply += '\n> default options:'
            reply += '\n> steps:30  (20-40)'
            reply += '\n> height:512  (512-1024)'
            reply += '\n> width:512  (512-1024)'
            reply += '\n> seed:number'
            reply += '\n> favourites: https://t.me/+UyMF3ArBe4pmYzQ5'

        if not reply:
            settings = {
                'steps': 30,
                'height': 512,
                'width': 512,
                'seed': random.randint(11111, 19999),
            }

            maximums = {
                'steps': 40,
                'height': 1024,
                'width': 1024,
            }

            minimums = {
                'steps': 20,
                'height': 512,
                'width': 512,
            }

            for match in re.findall('((\w+)[:=](\d+))', to_send):
                whole = match[0]
                setting = match[1]
                value = int(match[2])

                to_send = to_send.replace(whole, '').replace('  ', ' ')

                if setting in maximums and value > maximums[setting]:
                    reply = '> {} should be <= {}'.format(setting, maximums[setting])
                    break

                if setting in minimums and value < minimums[setting]:
                    reply = '> {} should be >= {}'.format(setting, minimums[setting])
                    break

                if setting in settings:
                    settings[setting] = value


        if not reply:
            try:
                to_send = to_send.replace('@', '')
                to_send = to_send.replace('!', '')
                to_send = to_send.strip()
                prompt = to_send

                if len(str(settings['seed'])) == 4:
                    enhancements1 = [
                        ', fantasy, intricate, elegant, digital painting, concept art, sharp focus, illustration by greg rutkowski, 4k',
                        ', synthwave style, artstation, detailed, award winning, dramatic lighting, miami vice, oil on canvas',
                        ', highly detailed painting by dustin nguyen, akihiko yoshida, greg tocchini, 4 k, trending on artstation, 8 k',
                        ', cinematic lighting, insanely detailed, intricate, artstation, cgsociety, painted by simon stalenhag, concept art, illustration, sharp focus',
                        ', photorealistic, highly detailed, illustration, lifelike, highly detailed, intricate, octane render, sharp focus, cyberpunk',
                        ', by ghibli, makoto shinkai, bright and transparent animation style, anime art wallpaper 4 k, trending on artstation',
                        ', painterly, sharp focus, illustration, art by lois royo',
                        ', in the style of gta cover art, by stephen bliss, trending on artstation, pixiv, 8 k',
                        ', d & d, fantasy concept art, octane render, 8 k',
                        ', peter mohrbacher, charlie bowater, artstation, craig mullins, cinematic, digital painting, deviantart, cinematic lighting, 4 k',
                        ', tristan eaton, victo ngai, maxfield parrish, artgerm, koons, ryden, intricated details',
                        ', by soul bass, kawase hasui, moebius and edward hopper, colorful flat surreal design, xray hd, 8 k',
                        ', manga cover art, detailed color portrait, artstation trending, 8 k',
                        ', a hacker hologram by penny patricia poppycock, pixabay contest winner, holography, irridescent, photoillustration, maximalist vaporwave - n 9',
                        ', breathtaking, epic, stunning, gorgeous, much detail, much wow, cgsociety, wlop, pixiv, behance, deviantart, masterpiece',
                        ', sharp focus, dynamic lighting, elegant, harmony, beauty, masterpiece, by roberto ferry, illustration, blue background',
                        ', ultradetailed environment, sharp focus, cinematic lighting, by alphonse maria mucha and kim jung gi',
                        ', concept art by jama jurabaev, cinematic shot, trending on artstation, high quality, brush stroke, hyperspace, vibrant colors',
                        ', hyper detailed, sharp focus, bokeh, unreal engine, ray tracing, cute, fantasy, sci fi, purple lights, tiny, small',
                        ', fancy, ultra detailed, brush strokes, digital painting, cinematic, yoshitaka amano, andy warhol, ultra realistic',
                        ', by tsuyoshi nagano, illustration, cinematic lighting, hyperdetailed, 8 k, symmetrical, frostbite 3 engine, cryengine, dof, trending on artstation, digital art, crepuscular ray',
                        ', retrowave, tonal separation, black background, neon, black, glitch, strong contrast, pinterest, trending on artstation',
                        ', basic background, krenz cushart, mucha, ghibli, by joaquin sorolla rhads leyendecker, by ohara koson',
                    ]

                    if ',' not in prompt:
                        prompt += enhancements1[settings['seed'] % len(enhancements1)]
                elif str(settings['seed'])[0] == '1':
                    pass

                start = time.time()
                dream_res = dreamer(prompt, **settings)
                finish = time.time()

                url_enc = dream_res['data'][0][0]
                b64_str = url_enc[22:]

                reply_file = base64.decodebytes(bytes(b64_str, 'utf-8'))

                latency = finish - start
                reply = '> generated "{}" for {} in {} s with seed:{} steps:{}'.format(
                    to_send,
                    name,
                    str(latency)[:3],
                    settings['seed'],
                    settings['steps']
                )

            except BaseException as e:
                logging.exception(e)
                reply = '> error'

        await event.reply(reply, file=reply_file, link_preview=False)


@client.on(events.MessageEdited)
async def message_edited(event):
    await message_handler(event)

@client.on(events.NewMessage)
async def new_message(event):
    await message_handler(event)


client.start()
client.run_until_disconnected()