import gradio as gr
import numpy as np
import cv2
import os
import requests
import sys
from PIL import Image

# ─── বাগ ফিক্স ────────────────────────────────────────────────────────────────
import torchvision.transforms.functional
sys.modules['torchvision.transforms.functional_tensor'] = torchvision.transforms.functional

# ─── মডেল সেটআপ ──────────────────────────────────────────────────────────────
MODEL_PATH = "RealESRGAN_x4plus.pth"

if not os.path.exists(MODEL_PATH):
    print("মডেল ডাউনলোড হচ্ছে...")
    r = requests.get("https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth", stream=True)
    with open(MODEL_PATH, "wb") as f:
        for chunk in r.iter_content(chunk_size=8192):
            f.write(chunk)

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

model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
upsampler = RealESRGANer(scale=4, model_path=MODEL_PATH, model=model, tile=128, tile_pad=10, pre_pad=0, half=False)

# ─── মূল কাজ (কোনো ফ্যান্সি ডিজাইন ছাড়া) ──────────────────────────────────────
def process_image(img):
    if img is None:
        return None, None
        
    # RGB থেকে BGR
    img_np = np.array(img.convert("RGB"))
    img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
    
    # 4K আপস্কেল
    output_bgr, _ = upsampler.enhance(img_bgr, outscale=4)
    
    # BGR থেকে RGB
    output_rgb = cv2.cvtColor(output_bgr, cv2.COLOR_BGR2RGB)
    out_img = Image.fromarray(output_rgb)
    
    # 4K সাইজ নিশ্চিত করা
    if out_img.size != (3840, 2160):
        out_img = out_img.resize((3840, 2160), Image.LANCZOS)
        
    # সেভ করা
    out_path = "/tmp/output.png"
    out_img.save(out_path, "PNG")
    
    return out_img, out_path

# ─── সিম্পল ইউজার ইন্টারফেস ──────────────────────────────────────────────────
with gr.Blocks() as demo:
    gr.Markdown("# 🚀 4K AI Upscaler (Basic Version)")
    gr.Markdown("সব ডিজাইন বাদ দেওয়া হয়েছে। শুধু ছবি আপলোড করুন এবং রেজাল্ট নিন।")
    
    with gr.Row():
        inp = gr.Image(type="pil", label="ছবি আপলোড করুন")
        out = gr.Image(type="pil", label="4K আউটপুট")
        
    btn = gr.Button("✨ 4K তে রূপান্তর করুন", variant="primary")
    file_out = gr.File(label="4K ছবি ডাউনলোড করুন")
    
    btn.click(fn=process_image, inputs=inp, outputs=[out, file_out])

if __name__ == "__main__":
    demo.queue().launch()