# 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="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="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()




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

# =====================================================
# Global Variables (Lazy Loading)
# =====================================================

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


# =====================================================
# Load AI Models Only Once
# =====================================================

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

        print("==========================================")
        print("Loading RealESRGAN Model...")
        print("==========================================")

        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
        )

        print("==========================================")
        print("Loading GFPGAN Model...")
        print("==========================================")

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

        print("==========================================")
        print("AI Models Loaded Successfully")
        print("==========================================")

    return _face_enhancer


# =====================================================
# Image Enhancement Service
# =====================================================

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()