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

import uuid
import time
import os
import requests
import gradio as gr
import simplejson as json
import oss2

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:
        raise gr.Error('登录错误，请确认你的720yun账号是否正确')


def create_panorama(prompt, negative_prompt):
    try:
        data = {
            'api_key': os.environ['a2'],
            'generator': 'stable-skybox-trt',
            'prompt': prompt[0:598],
            'negative_text': negative_prompt[0:398]
        }
        resp = requests.post(os.environ['a1'], data=data)
        print(resp.text)
        if resp.status_code == 200:
            time.sleep(20)
            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(0.5)
    except Exception:
        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:
        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, aerial view, epic composition, post apocalyptic, sci-fi, futuristic, fantasy, by Jan Urschel and Sergey Vasnev and Emmanuel Shiu and Michal Karcz, ornate, cinematic, cinematic lighting, light effect, epic, octane render, unreal engine, 4k hd wallpaper, trending on artstation and cgsociety
    """)
    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('Run')
    btn.click(fn=main, inputs=[login_id, password, prompt, negative_prompt, session], outputs=out, show_progress=True)

demo.queue().launch()