import os
import cv2
import torch
import glob
import requests
import numpy as np
import gradio as gr
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer

# ==========================================
# 1. AUTO-PATCHER (Bedah Kode Otomatis)
# ==========================================
# Mengatasi error 'torchvision' yang sering terjadi di server baru
def patch_basicsr():
    try:
        import basicsr
        pkg_path = os.path.dirname(basicsr.__file__)
        target_file = os.path.join(pkg_path, 'data', 'degradations.py')
        
        if os.path.exists(target_file):
            with open(target_file, 'r') as f: data = f.read()
            
            old_str = 'from torchvision.transforms.functional_tensor import rgb_to_grayscale'
            new_str = 'from torchvision.transforms.functional import rgb_to_grayscale'
            
            if old_str in data:
                data = data.replace(old_str, new_str)
                with open(target_file, 'w') as f: f.write(data)
                print("✅ Auto-Patch: Library basicsr berhasil diperbaiki.")
    except Exception as e:
        print(f"⚠️ Auto-Patch Warning: {e}")

patch_basicsr()

# ==========================================
# 2. PERSIAPAN MODEL
# ==========================================
# Cek ketersediaan GPU (Hugging Face Free Tier biasanya CPU, Paid dapat GPU)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"🚀 Running on: {device}")

# Download Model Pre-trained jika belum ada
model_url = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth'
model_path = 'RealESRGAN_x4plus.pth'

if not os.path.exists(model_path):
    print("⏳ Mendownload model AI...")
    response = requests.get(model_url)
    with open(model_path, 'wb') as f:
        f.write(response.content)
    print("✅ Model siap!")

# Load Model ke Memory
def load_model():
    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
    netscale = 4
    
    upsampler = RealESRGANer(
        scale=netscale,
        model_path=model_path,
        dni_weight=None,
        model=model,
        tile=200,       # Tile kecil agar aman untuk CPU/RAM kecil
        tile_pad=10,
        pre_pad=0,
        half=False,     # False agar aman di CPU (fp32)
        device=device
    )
    return upsampler

# Inisialisasi Upsampler Global
upsampler = load_model()

# ==========================================
# 3. LOGIKA UTAMA (PIPELINE)
# ==========================================
def process_image(img, scale, face_enhance):
    if img is None:
        return None

    # --- MODE DOWNSCALE (MENGECILKAN) ---
    if scale < 1.0:
        h, w = img.shape[:2]
        new_w, new_h = int(w * scale), int(h * scale)
        # Gunakan INTER_AREA untuk hasil pengecilan terbaik (tajam)
        output = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
        return output

    # --- MODE UPSCALE (MEMBESARKAN) ---
    else:
        try:
            # Upscale menggunakan Real-ESRGAN
            # outscale=scale akan otomatis resize hasil akhir sesuai keinginan user
            output, _ = upsampler.enhance(img, outscale=scale)
            
            # Fitur Perbaikan Wajah (Optional)
            if face_enhance:
                from gfpgan import GFPGANer
                face_enhancer = GFPGANer(
                    model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',
                    upscale=scale,
                    arch='clean',
                    channel_multiplier=2,
                    bg_upsampler=upsampler
                )
                _, _, output = face_enhancer.enhance(output, has_aligned=False, only_center_face=False, paste_back=True)
            
            return output
            
        except RuntimeError as e:
            # Tangani error jika RAM tidak cukup
            if 'out of memory' in str(e):
                return None
            raise e

# ==========================================
# 4. ANTARMUKA GRADIO
# ==========================================
title = "Smart Image Resizer & Upscaler"
description = """
**Fitur Utama:**
1. **Upscale Cerdas:** Menggunakan Real-ESRGAN untuk memperbesar gambar tanpa pecah.
2. **Downscale Tajam:** Menggunakan algoritma resampling area untuk memperkecil gambar dengan detail tinggi.
3. **Face Restore:** Memperbaiki wajah yang buram (opsional).

*Catatan: Jika berjalan di CPU (HF Free Tier), proses upscale gambar besar mungkin memakan waktu 30-60 detik.*
"""

iface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(label="Upload Gambar", type="numpy"),
        gr.Slider(0.1, 4.0, value=2.0, step=0.1, label="Skala Resolusi (0.5 = Kecil, 2.0 = Besar)"),
        gr.Checkbox(label="Perbaiki Wajah (Face Restoration)", value=False)
    ],
    outputs=gr.Image(label="Hasil Proses"),
    title=title,
    description=description,
    allow_flagging="never"
)

if __name__ == "__main__":
    iface.launch()