import cv2
import numpy as np
import torch

from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer
from gfpgan import GFPGANer

device = "cuda" if torch.cuda.is_available() else "cpu"

# ---------------------------
# RealESRGAN
# ---------------------------

rrdbnet = RRDBNet(
    num_in_ch=3,
    num_out_ch=3,
    num_feat=64,
    num_block=23,
    num_grow_ch=32,
    scale=4
)

bg_upsampler = RealESRGANer(
    scale=4,
    model_path=r"app\models\RealESRGAN_x4plus.pth",
    model=rrdbnet,
    tile=256,
    tile_pad=10,
    pre_pad=0,
    half=torch.cuda.is_available()
)

# ---------------------------
# GFPGAN
# ---------------------------

face_enhancer = GFPGANer(
    model_path=r"app\models\GFPGANv1.4.pth",
    upscale=4,
    arch="clean",
    channel_multiplier=2,
    bg_upsampler=bg_upsampler
)

# ---------------------------
# Service
# ---------------------------

def enhance_image_service(image_bytes: bytes) -> bytes:

    nparr = np.frombuffer(image_bytes, np.uint8)

    img = cv2.imdecode(
        nparr,
        cv2.IMREAD_COLOR
    )

    if img is None:
        raise Exception("Invalid image")

    print("Input:", img.shape)

    cropped_faces, restored_faces, output = face_enhancer.enhance(
        img,
        has_aligned=False,
        only_center_face=False,
        paste_back=True
    )

    print("Output:", output.shape)

    success, encoded_img = cv2.imencode(
        ".png",
        output
    )

    if not success:
        raise Exception("Failed to encode image")

    return encoded_img.tobytes()