# import cv2
# import numpy as np
# import torch
# import threading

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

# _face_enhancer = None
# _bg_upsampler = None
# _model_lock = threading.Lock()


# def get_face_enhancer():
#     global _face_enhancer
#     global _bg_upsampler

#     # Already loaded
#     if _face_enhancer is not None:
#         return _face_enhancer

#     # Thread Safe
#     with _model_lock:

#         # Double Check
#         if _face_enhancer is not None:
#             return _face_enhancer

       
#         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="app/models/RealESRGAN_x4plus.pth",
#             model=rrdbnet,

#             # CPU Server Optimization
#             tile=64,
#             tile_pad=10,
#             pre_pad=0,

#             half=False
#         )

#         _face_enhancer = GFPGANer(
#             model_path="app/models/GFPGANv1.4.pth",
#             upscale=4,
#             arch="clean",
#             channel_multiplier=2,
#             bg_upsampler=_bg_upsampler
#         )

       
#     return _face_enhancer



# def enhance_image_service(image_bytes: bytes) -> bytes:

#     enhancer = get_face_enhancer()

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

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

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

#     print(f"Input Image Shape : {img.shape}")

#     _, _, output = enhancer.enhance(
#         img,
#         has_aligned=False,
#         only_center_face=False,
#         paste_back=True
#     )

#     print(f"Output Image Shape : {output.shape}")

#     success, encoded = cv2.imencode(
#         ".png",
#         output
#     )

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

#     return encoded.tobytes()





import gc
import cv2
import numpy as np

from app.modules.enhanceImage.loader import (
    get_face_enhancer
)


def enhance_image_service(image_bytes: bytes) -> bytes:

    try:

        # =====================================
        # Decode Image
        # =====================================

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

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

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

        print("----------------------------------")
        print("Input Shape :", img.shape)

        enhancer = get_face_enhancer()

        # =====================================
        # Face Enhancement Only
        # =====================================

        cropped_faces, restored_faces, output = enhancer.enhance(

            img,

            has_aligned=False,

            only_center_face=False,

            paste_back=True

        )

        print("Detected Faces :", len(restored_faces))

        if output is None:
            output = img

        print("Output Shape :", output.shape)

        # =====================================
        # Encode PNG
        # =====================================

        success, encoded = cv2.imencode(

            ".png",

            output,

            [

                cv2.IMWRITE_PNG_COMPRESSION,

                3

            ]

        )

        if not success:

            raise Exception(

                "Failed to encode image"

            )

        print("----------------------------------")

        return encoded.tobytes()

    finally:

        gc.collect()