import gradio as gr
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, models
from art.attacks.evasion import FastGradientMethod
from art.estimators.classification import PyTorchClassifier
from PIL import Image
import numpy as np
import os
import io
from blind_watermark import WaterMark

# Pretrained ResNet50 모델 불러오기 (ImageNet 사전 훈련)
model = models.resnet50(pretrained=True)

# CIFAR-10에 맞춰 마지막 분류 레이어 수정
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, 10)

# 모델을 GPU로 이동
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# 손실 함수와 옵티마이저 설정
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# PyTorchClassifier 생성
classifier = PyTorchClassifier(
    model=model,
    loss=criterion,
    optimizer=optimizer,
    input_shape=(3, 64, 64),
    nb_classes=10,
)

# 이미지 전처리 함수
def preprocess_image(image):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])
    return transform(image).unsqueeze(0).to(device)

# FGSM 공격 적용 및 이미지 처리 함수
def generate_adversarial_image(image, eps_value):
    img_tensor = preprocess_image(image)

    # FGSM 공격 설정
    attack = FastGradientMethod(estimator=classifier, eps=eps_value)

    # 적대적 예제 생성
    adv_img_tensor = attack.generate(x=img_tensor.cpu().numpy())
    adv_img_tensor = torch.tensor(adv_img_tensor).to(device)

    # 적대적 이미지 변환
    adv_img_np = adv_img_tensor.squeeze(0).cpu().numpy()
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    adv_img_np = (adv_img_np * std[:, None, None]) + mean[:, None, None]
    adv_img_np = np.clip(adv_img_np, 0, 1)
    adv_img_np = adv_img_np.transpose(1, 2, 0)

    # PIL 이미지로 변환
    adv_image_pil = Image.fromarray((adv_img_np * 255).astype(np.uint8))
    
    return adv_image_pil

# 워터마크 삽입 함수
def apply_watermark(image_pil, wm_text="텍스트 삽입", password_img=000, password_wm=000):
    bwm = WaterMark(password_img=password_img, password_wm=password_wm)

    # 이미지 바이트 데이터를 임시 파일로 저장
    temp_image_path = "temp_image.png"
    image_pil.save(temp_image_path)

    # temp_image_path 경로로 워터마크 삽입 처리
    bwm.read_img(temp_image_path)
    bwm.read_wm(wm_text, mode='str')
    
    # 워터마크 삽입
    output_path = "watermarked_image.png"
    bwm.embed(output_path)

    # 삽입된 워터마크 이미지 파일을 다시 읽어서 PIL 이미지로 변환
    result_image = Image.open(output_path)

    # 임시 파일 삭제
    os.remove(temp_image_path)
    os.remove(output_path)

    return result_image

# 전체 이미지 처리 함수
def process_image(image, eps_value, wm_text, password_img, password_wm):
    # 적대적 이미지 생성
    adv_image = generate_adversarial_image(image, eps_value)
    
    # 적대적 이미지에 워터마크 삽입
    watermarked_image = apply_watermark(adv_image, wm_text, int(password_img), int(password_wm))

    return watermarked_image

# Gradio 인터페이스 정의
gr.Interface(
    fn=process_image,
    inputs=[gr.Image(type="pil", label="이미지를 업로드하세요"),  # 이미지 업로드 필드
        gr.Slider(0.1, 1.0, step=0.1, value=0.3, label="Epsilon 값 설정 (노이즈 강도)"),  # epsilon 값 슬라이더
        gr.Textbox(label="워터마크 텍스트 입력", value="텍스트 삽입"),  # 워터마크 텍스트 입력 필드
        gr.Number(label="이미지 비밀번호", value=000),  # 이미지 비밀번호 입력 필드
        gr.Number(label="워터마크 비밀번호", value=000)  # 워터마크 비밀번호 입력 필드
    ],
    outputs=gr.Image(type="pil", label="워터마크가 삽입된 이미지")  # 결과 이미지 출력 필드
).launch()