import cv2
from basicsr.archs.rrdbnet_arch import RRDBNet
from web.log import logger
from realesrgan import RealESRGANer
from gfpgan import GFPGANer
import numpy as np

# restorer
def create_real_esrganer(model_path, num_block, scale=4):
    upsampler = RealESRGANer(
        scale=scale,
        model_path=model_path,
        dni_weight=None,
        model=RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=num_block, num_grow_ch=32, scale=scale),
        tile=0,
        tile_pad=10,
        pre_pad=0,
        half=False,
        gpu_id=None)

    face_enhancer = GFPGANer(
        model_path='weights/GFPGANv1.4.pth',
        upscale=scale,
        arch='clean',
        channel_multiplier=2,
        bg_upsampler=upsampler)
    return upsampler, face_enhancer

def convert_img(path: str, origin: str, output_name: str, face_enhance_open:bool, amine:bool):
    scale = 4
    if amine:
        upsampler, face_enhancer = create_real_esrganer('weights/RealESRGAN_x4plus_anime_6B.pth', 6, scale)
    else:
        upsampler, face_enhancer = create_real_esrganer('weights/RealESRGAN_x4plus.pth', 23, scale)
    img = cv2.imread(f"{path}/{origin}", cv2.IMREAD_UNCHANGED)
    logger.info("convert_img img shape: %s", img.shape)
    try:
        if face_enhance_open:
            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
        else:
            output, _ = upsampler.enhance(img, outscale=scale)
    except Exception as error:
        logger.info('Error %s', error)
        logger.info('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
    else:
        cv2.imwrite(f"{path}/{output_name}", output)


def convert_img_row(data: bytes, output_name: str, face_enhance_open:bool, amine:bool):
    if amine:
        upsampler, face_enhancer = create_real_esrganer('weights/RealESRGAN_x4plus_anime_6B.pth', 6)
    else:
        upsampler, face_enhancer = create_real_esrganer('weights/RealESRGAN_x4plus.pth', 23)
    img = cv2.imdecode(np.fromstring(data, np.uint8), cv2.IMREAD_UNCHANGED)
    try:
        if face_enhance_open:
            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
        else:
            output, _ = upsampler.enhance(img, outscale=4)
    except RuntimeError as error:
        print('Error', error)
        print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
    else:
        return cv2.imencode(f'.{output_name}', output)[1].tobytes()