import subprocess
subprocess.run("pip install opencv-python os gc torch torchaudio torchvision gfpgan basicsr realesrgan")
import os
import cv2
import time
import gc
import torch
from gfpgan import GFPGANer
from basicsr.archs.rrdbnet_arch import RRDBNet
from basicsr.utils.download_util import load_file_from_url
from realesrgan import RealESRGANer

# === GPU MEMORY MONITORING ===
def print_simple_gpu_memory():
    allocated = torch.cuda.memory_allocated(0) / 1024**3
    reserved = torch.cuda.memory_reserved(0) / 1024**3
    print(f"[GPU Memory] Allocated: {allocated:.2f} GB | Reserved: {reserved:.2f} GB")

# === GFPGAN STEP ===
def run_gfpgan(image_path, output_path, model_path='GFPGAN\GFPGAN\experiments\pretrained_models\GFPGANv1.4.pth'):
    gfpganer = GFPGANer(
        model_path=model_path,
        upscale=2,
        arch='clean',
        channel_multiplier=2,
        bg_upsampler=None
    )

    img = cv2.imread(image_path)
    if img is None:
        raise FileNotFoundError(f"Input image not found: {image_path}")

    _, _, restored_img = gfpganer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
    cv2.imwrite(output_path, restored_img)
    print(f"[+] GFPGAN output saved: {output_path}")
    return output_path

# === RealESRGAN STEP ===
def run_realesrgan(input_path, output_path, model_name='RealESRGAN_x4plus', outscale=4, gpu_id=None):
    if model_name == 'RealESRGAN_x4plus':
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
        netscale = 4
        file_url = ['https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth']
    else:
        raise NotImplementedError(f'Model {model_name} not implemented')

    model_path = os.path.join('weights', model_name + '.pth')
    if not os.path.isfile(model_path):
        ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
        for url in file_url:
            model_path = load_file_from_url(url=url, model_dir=os.path.join(ROOT_DIR, 'weights'), progress=True, file_name=None)

    upsampler = RealESRGANer(
        scale=netscale,
        model_path=model_path,
        dni_weight=None,
        model=model,
        tile=512,  # <<< IMPORTANT for 4 GB VRAM!
        tile_pad=10,
        pre_pad=0,
        half=False,
        gpu_id=gpu_id
    )

    img = cv2.imread(input_path, cv2.IMREAD_UNCHANGED)
    if img is None:
        raise FileNotFoundError(f'Input image not found: {input_path}')

    output, _ = upsampler.enhance(img, outscale=outscale)

    output_dir = os.path.dirname(output_path)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir, exist_ok=True)

    cv2.imwrite(output_path, output)
    print(f"[+] RealESRGAN output saved: {output_path}")
    return output_path

# === Combined Workflow ===
def combined_enhance(input_img_path, rescale_factor, output_dir):
    start_time = time.perf_counter()

    base_name = os.path.splitext(os.path.basename(input_img_path))[0]

    # Step 1: GFPGAN enhancement
    gfpgan_out = os.path.join(output_dir, f"{base_name}_GFPGAN.png")
    run_gfpgan(input_img_path, gfpgan_out)

    print("[GPU Memory after GFPGAN]")
    print_simple_gpu_memory()

    # Free memory before loading RealESRGAN
    gc.collect()
    torch.cuda.empty_cache()
    print("[*] Cleared GPU cache before RealESRGAN")

    # Step 2: RealESRGAN enhancement using GFPGAN result
    realesrgan_out = os.path.join(output_dir, f"{base_name}_GFPGAN_RealESRGAN_combined.png")
    run_realesrgan(gfpgan_out, realesrgan_out, outscale=rescale_factor)

    print("[GPU Memory after RealESRGAN]")
    print_simple_gpu_memory()

    print(f"[***] Final combined image saved at: {realesrgan_out}")

    # Get resolution and file size info
    img = cv2.imread(realesrgan_out)
    height, width = img.shape[:2]
    file_size_bytes = os.path.getsize(realesrgan_out)
    file_size_mb = file_size_bytes / (1024 * 1024)

    end_time = time.perf_counter()
    elapsed = end_time - start_time
    mins, secs = divmod(elapsed, 60)
    hours, mins = divmod(mins, 60)

    print(f"[***] Image resolution: {width}x{height}, file size: {file_size_mb:.2f} MB")
    print(f"[***] Total processing time: {int(hours):02d}h:{int(mins):02d}m:{secs:05.2f}s")

    return realesrgan_out

import gradio as gr
import os

# Your core functions like combined_enhance must be defined/imported before this
def wrapped_combined_enhance(image, scale_factor):
    temp_input_path = "temp_input.png"
    temp_output_dir = "output"
    os.makedirs(temp_output_dir, exist_ok=True)
    image.save(temp_input_path)

    try:
        final_path = combined_enhance(temp_input_path, int(scale_factor), temp_output_dir)
        return final_path, "Enhancement succeeded!"
    except Exception as e:
        return None, f"[ERROR]: {str(e)}"

with gr.Blocks(title="AI Face Enhancer (GFPGAN + RealESRGAN)") as demo:
    gr.Markdown("# ✨ Face Enhancer Pro")

    with gr.Row():
        with gr.Column():
            input_img   = gr.Image(type="pil", label="Upload Image")
            scale_slider= gr.Slider(1, 4, value=2, step=1, label="Upscale Factor")
            enhance_btn = gr.Button("Enhance")
            status_box  = gr.Textbox(label="Status / Logs", lines=6, interactive=False)
        with gr.Column():
            output_img  = gr.Image(label="Enhanced Output")

    def ui_callback(image, scale):
        if image is None:
            return None, "⚠️ Please upload an image first."
        out_path, status = wrapped_combined_enhance(image, scale)
        # return filepath (or None) and status string
        return out_path, status

    enhance_btn.click(
        fn=ui_callback,
        inputs=[input_img, scale_slider],
        outputs=[output_img, status_box]
    )

demo.launch()