from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Optional
import api_module as mj
import random
from io import BytesIO
import openai
from transparent_background import Remover
import requests
from PIL import Image
from qiniu import Auth, put_data, etag, Zone
import replicate
import os


mj_api_key = os.environ['MJ_API_KEY']
openai.api_key = os.environ['OPENAI_API_KEY'] 
access_key = os.environ['ACCESS_KEY']
secret_key = os.environ['SECRET_KEY']
bucket_name = 'mystory-games'



def translate_to_english(text):
    response = openai.ChatCompletion.create(
        model="gpt-3.5",
        messages=[
            {
                "role": "system",
                "content": "accurately translating Chinese text to English, while keeping the context and idiomatically correct in English, ensure that only translated content is replied to \n"
            },
            {
                "role": "user",
                "content": text
            },
        ],
        temperature=0.2,
        max_tokens=256,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0
    )

    translated_text = response['choices'][0]['message']['content'].strip()
    return translated_text

def translate_to_chinese(text):
    response = openai.ChatCompletion.create(
        model="gpt-4",
        messages=[
            {
                "role": "system",
                "content": "accurately translating English text to Chinese, while keeping the context and idiomatically correct in Chinese, ensure that only translated content is replied to  \n"
            },
            {
                "role": "user",
                "content": text
            },
        ],
        temperature=0.2,
        max_tokens=256,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0
    )

    translated_text = response['choices'][0]['message']['content'].strip()
    return translated_text

def avatar_create(text):
    #translate to english
    text = translate_to_english(text)
    #add text to prompt
    prompt = f'{text}, concept art by Pixar, cgsociety, character, furry art, full body shot, white background, Realistic 3d render --v 5.2'
    #make imagine request
    max_retries = 5
    retries = 0
    while retries < max_retries:
        task_id = mj.make_imagine_request(mj_api_key, prompt)
        if task_id is not None:
            break
        else:
            retries += 1
            print(f'Failed to get task id, retrying {retries}/{max_retries}...')
    if retries == max_retries:
        print('Failed to get task id, aborting')
        return None
    
    #fetch request
    # make sure the imagine process has finished
    while True:
        image_url = mj.fetch_request(mj_api_key, task_id)
        if image_url is not None:
            break
        elif image_url is None:
            print('Imagine process failed, exiting...')
            raise SystemExit('Imagine process failed')
    return task_id

def avatar_remove_background(image_url):
    remover = Remover() # default settings
    # downlaod image
    response = requests.get(image_url)
    # convert to PIL image
    img = Image.open(BytesIO(response.content)).convert('RGB')
    # remove background
    out = remover.process(img, type='rgba')
    image_io = BytesIO()
    Image.fromarray(out).save(image_io, format='PNG')
    image_data = image_io.getvalue()

    # Upload the image data
    q = Auth(access_key, secret_key)
    token = q.upload_token(bucket_name)
    ret, info = put_data(token, None, image_data, mime_type="image/png")
    if info.status_code == 200:
        cdn_url = 'http://image1.juramaia.com/' + ret['key']
        return cdn_url
    else:
        print('Failed to upload image to CDN')
        return None

def avatar_upscale(task_id, index):
    task_id = mj.make_upscale_request(mj_api_key, task_id, index)
    # make sure the upscale task sucessfully
    if task_id is not None:
        pass
    elif task_id is None:
        print('Upscale process failed, exiting...')
        raise SystemExit('Upscale process failed')
    # make sure the upscale process has finished
    while True:
        upscale_image_url = mj.fetch_request(mj_api_key, task_id)
        if upscale_image_url is not None:
            break
        elif upscale_image_url is None:
            print('Upscale process failed, exiting...')
            raise SystemExit('Upscale process failed')
    return upscale_image_url

def main_process(text):
    #create avatar
    task_id = avatar_create(text)
    print(f'\navatar create successful task_id is {task_id}\n')
    #upscale avatar and remove background
    index = random.randint(1, 4)
    while True:
        print(f'upsclae index is {index} and task_id is {task_id}')
        upscale_image_url = avatar_upscale(task_id, index)
        if upscale_image_url is not None:
            break
        elif upscale_image_url is None:
            print('Upscale process failed, exiting...')
            raise SystemExit('Upscale process failed')
    print(f'upscale_image_url is {upscale_image_url}')    
    return upscale_image_url

def analyze_process(url):
    print(f'\nanalyzing image: {url}')
    eng_output = replicate.run(
        "andreasjansson/blip-2:4b32258c42e9efd4288bb9910bc532a69727f9acd26aa08e175713a0a857a608",
        input={
            "image": url, 
            "caption": True,
            "temperature": 0.5
        }
    )
    output = translate_to_chinese(eng_output)
    print(f'analyzed caption: {output}\n')
    return output


class Item(BaseModel):
    text: str

api = FastAPI()

@api.post("/main_process/")
async def main_api(item: Item):
    try:
        image_url = main_process(item.text)
        return {"image_url": image_url}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))
    
@api.post("/analyze_process/")
async def analyze_api(item: Item):
    try:
        output = analyze_process(item.text)
        return {"output": output}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))