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

class EnhancementModels:
    """Lazy load models to save memory until needed."""
    _x2_upscaler = None
    _x4_upscaler = None
    _face_restorer = None

    @classmethod
    def get_x2_upscaler(cls):
        if cls._x2_upscaler is None:
            print("Loading RealESRGAN x2...")
            model_path = "/data/RealESRGAN_x2.pth"
            if not os.path.exists(model_path):
                from realesrgan.utils import download_weights
                download_weights('RealESRGAN_x2', model_path)
            model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
            cls._x2_upscaler = RealESRGANer(
                scale=2, model_path=model_path, model=model,
                tile=0, tile_pad=10, pre_pad=0, half=False
            )
        return cls._x2_upscaler

    @classmethod
    def get_x4_upscaler(cls):
        if cls._x4_upscaler is None:
            print("Loading RealESRGAN x4 (this may take a moment)...")
            model_path = "/data/RealESRGAN_x4plus.pth"
            if not os.path.exists(model_path):
                from realesrgan.utils import download_weights
                download_weights('RealESRGAN_x4plus', model_path)
            model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
            cls._x4_upscaler = RealESRGANer(
                scale=4, model_path=model_path, model=model,
                tile=0, tile_pad=10, pre_pad=0, half=False
            )
        return cls._x4_upscaler

    @classmethod
    def get_face_restorer(cls):
        if cls._face_restorer is None:
            print("Loading GFPGAN for face restoration...")
            cls._face_restorer = GFPGANer(
                model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
                upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None
            )
        return cls._face_restorer

def upscale_image(image_bytes: bytes, scale: int = 2, restore_face: bool = False) -> bytes:
    import cv2
    import numpy as np
    nparr = np.frombuffer(image_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    if scale == 2:
        upscaler = EnhancementModels.get_x2_upscaler()
    elif scale == 4:
        upscaler = EnhancementModels.get_x4_upscaler()
    else:
        raise ValueError("Only scale 2 or 4 supported")

    output, _ = upscaler.enhance(img, outscale=scale)
    
    if restore_face:
        face_restorer = EnhancementModels.get_face_restorer()
        _, _, output = face_restorer.enhance(output, has_aligned=False, only_center_face=False, paste_back=True)
    
    _, encoded = cv2.imencode('.png', output)
    return encoded.tobytes()