#!/usr/local/bin/python3
# -*- coding: utf-8 -*-

import uuid
import time
import os
import logging
import logging.handlers
import requests
import gradio as gr
import simplejson as json
import oss2

LOGGER = logging.getLogger('skybox')
HANDLER = logging.handlers.RotatingFileHandler(
    './logs/skybox.log', mode='a', maxBytes=20 * 1000 * 1000, backupCount=5, encoding='utf-8')
FORMATTER = logging.Formatter('%(asctime)s - %(funcName)s - %(lineno)d - %(levelname)s - %(message)s')
HANDLER.setFormatter(FORMATTER)
LOGGER.propagate = 0
LOGGER.addHandler(HANDLER)
LOGGER.setLevel(logging.DEBUG)

CSS = '.gradio-container a {color:#b7adf4 !important}'
HEADERS = {'app-key': os.environ['a3'], 'origin': os.environ['a4'], 'referer': os.environ['a4']}


def login(loginId, password):
    try:
        data = {'loginId': loginId, 'password': password}
        cookies = {'720yun_v8_session': str(uuid.uuid4())}

        resp = requests.post(os.environ['a5'], data=data, headers=HEADERS, cookies=cookies)
        if resp.status_code == 200:
            return resp.json()['data']['token']
    except Exception as e:
        LOGGER.error(e)
        raise gr.Error('登录错误，请确认你的720yun账号是否正确')


def create_panorama(prompt, negative_prompt):
    try:
        data = {
            'api_key': os.environ['a2'],
            'generator': 'stable-skybox',
            'prompt': prompt.strip()[0:598],
            'negative_text': negative_prompt.strip()[0:398]
        }
        resp = requests.post(os.environ['a1'], data=data)
        LOGGER.info(resp.text)
        if resp.status_code == 200:
            time.sleep(15)
            request_id = resp.json()['request']['id']
            flag = True
            countdown = 50
            while flag:
                countdown -= 1
                if countdown <= 0:
                    raise gr.Error('AI老师罢工了，请稍后重试')

                resp = requests.get(f'{os.environ["a1"]}/{request_id}?api_key={os.environ["a2"]}')
                if resp.status_code == 200:
                    progress_data = resp.json()['request']
                    if progress_data['progress'] == 100 and progress_data['status'] == 'complete':
                        flag = False
                        return request_id, progress_data['file_url']
                time.sleep(1)
    except Exception as e:
        LOGGER.error(e)
        raise gr.Error('AI老师罢工了，请稍后重试')


def upload(token, file_id, file_url):
    now = int(time.time())
    data = [{
        'fileId': f'{file_id}-{now}',
        'name': f'AI素材-{now}',
        'watermarked': 0,
        'size': 1024 * 1024 * 6,
        'albumId': 0,
        'exif': {},
        'gps': {},
        'panoId': 0,
        'action': 1
    }]
    HEADERS['app-authorization'] = token

    try:
        resp = requests.post(os.environ['a6'], data={'panos': json.dumps(data)}, headers=HEADERS)
        if resp.status_code == 200:
            resp_data = resp.json()['data'][0]
            access_key_id = resp_data['accessKeyId']
            security_token = resp_data['securityToken']
            accessKey_secret = resp_data['accessKeySecret']
            bucket_name = resp_data['bucketName']
            path = resp_data['path'][1:]
            endpoint = resp_data['endpointO']
            pano_id = resp_data['panoId']
            task_id = resp_data['taskId']
            expired = resp_data['expired']

            auth = oss2.StsAuth(access_key_id, accessKey_secret, security_token)
            bucket = oss2.Bucket(auth, endpoint, bucket_name)

            input_stream = requests.get(file_url)
            result = bucket.put_object(f'{path}/{pano_id}.jpg', input_stream)
            if result.status == 200:
                resp = requests.post(f'{os.environ["a6"]}/{task_id}', data={'status': 3, 'expired': expired}, headers=HEADERS)
                if resp.status_code == 200:
                    time.sleep(5)
                    flag = True
                    while flag:
                        pano_ids = [pano['id'] for pano in requests.get(os.environ["a7"], headers=HEADERS).json()['data']]
                        if pano_id not in pano_ids:
                            flag = False
                            break
                        time.sleep(0.5)
                    data = {
                        'name': f'AI全景作品-{now}',
                        'materials': json.dumps([{
                            'type': 1,
                            'id': pano_id
                        }]),
                        'templateId': 2,
                        'publishPlatform': 1,
                        'keywords': 'AI全景',
                        'source': 99
                    }
                    resp = requests.post(os.environ['a8'], data=data, headers=HEADERS)
                    if resp.status_code == 200:
                        return resp.json()['data']['tid']
    except Exception as e:
        LOGGER.error(e)
        raise gr.Error('AI老师罢工了，请稍后重试')


def main(loginId, password, prompt, negative_prompt, state, progress=gr.Progress()):
    try:
        if 'token' not in state:
            state['token'] = login(loginId, password)
        token = state['token']

        file_id, image = create_panorama(prompt, negative_prompt)
        panorama_id = upload(token, file_id, image)
        return f'https://www.720yun.com/vr/{panorama_id}'
    except Exception as e:
        return f'{e}'


with gr.Blocks(css=CSS) as demo:
    session = gr.State({})

    gr.Markdown("""
        # 创造属于你自己的AI全景
        1. 需要你的 **[720yun.com](https://www.720yun.com)** 账号.
        2. 由 **[720yun.com](https://www.720yun.com)** 提供支持.
        3. 描述例子: a beautiful matte painting, northernmost continent, a gigantic square fortress covered by blizzard, epic composition, post apocalyptic, sci-fi, futuristic, fantasy, by Jan Urschel and Sergey Vasnev and Emmanuel Shiu and Michal Karcz, cinematic, cinematic lighting, light effect, epic, octane render, unreal engine
    """)
    with gr.Row():
        with gr.Column():
            login_id = gr.Textbox(label='账号: (必填)', placeholder='720yun 账号')
        with gr.Column():
            password = gr.Textbox(label='密码: (必填)', type='password', placeholder='720yun 密码')
    with gr.Row():
        prompt = gr.Textbox(label='描述: (必填)', lines=4, placeholder='请使用英文输入描述，最多600个字符')
    with gr.Row():
        negative_prompt = gr.Textbox(label='负面描述', lines=2, placeholder='请使用英文输入负面描述，最多400个字符')
    with gr.Row():
        out = gr.Textbox(label='输出')

    btn = gr.Button('运行')
    btn.click(fn=main, inputs=[login_id, password, prompt, negative_prompt, session], outputs=out, show_progress=True)

demo.queue(concurrency_count=8).launch()