Kan1aSorata.AIPaint/AIPaint/DrawViewController.swift

//
//  DrawViewController.swift
//  AIPaint
//
//  Created by Fengyu He on 2022/11/29.
//

import UIKit
import Photos
import SnapKit
import SwiftyJSON
import SwiftUI
import CoreData

class DrawViewController: UIViewController {
    
    lazy var generateView: UIView = {
        let view = UIView()
        view.backgroundColor = .clear
        view.layer.borderWidth = 1
        view.layer.borderColor = UIColor.black.cgColor
        return view
    }()
    
    lazy var promptTextField: UITextField = {
        let textField = UITextField()
        textField.placeholder = "请输入你的作品的组成元素，以空格分割。"
        textField.layer.borderWidth = 1
        textField.layer.borderColor = UIColor.black.cgColor
        return textField
    }()
    
    lazy var negativePromptTextField: UITextField = {
        let textField = UITextField()
        textField.placeholder = "请输入你不想在作品中看到的元素，以空格分割。"
        textField.layer.borderWidth = 1
        textField.layer.borderColor = UIColor.black.cgColor
        return textField
    }()
    
    lazy var generateButton: UIButton = {
        let button = UIButton()
        button.layer.backgroundColor = UIColor.blue.cgColor
        button.addTarget(self, action: #selector(generateAction), for: .touchUpInside)
        button.setTitle("生成图片", for: .normal)
        button.isUserInteractionEnabled = true
        return button
    }()
    
    var longPress: UILongPressGestureRecognizer?
    
    lazy var generateImage: UIImageView = {
        let imageView = UIImageView()
        imageView.layer.borderWidth = 1
        imageView.layer.borderColor = UIColor.black.cgColor
        imageView.isUserInteractionEnabled = true
        return imageView
    }()
    
    var prompt = ""
    var tranPrompt = ""
    var negativePrompt = ""
    var tranNegativePrompt = ""
    var sessionHash = ""
    var delegate: AppDelegate?
    var context: NSManagedObjectContext?
    var painting: Painting?
    let alertController = UIAlertController(title: "保存成功!", message: nil, preferredStyle: .alert)
    
    override func viewDidLoad() {
        super.viewDidLoad()
        
        view.backgroundColor = .white
        self.title = "AI Paint"
        
        delegate = UIApplication.shared.delegate as? AppDelegate
        context = delegate!.persistentContainer.viewContext
        painting = NSEntityDescription.insertNewObject(forEntityName: "Painting", into: context!) as? Painting

        view.addSubview(generateView)
        generateView.addSubview(promptTextField)
        generateView.addSubview(negativePromptTextField)
        generateView.addSubview(generateButton)
        
        view.addSubview(generateImage)
        longPress = UILongPressGestureRecognizer(target: self, action: #selector(saveImageToPhotos))
        generateImage.addGestureRecognizer(longPress!)
        
        generateView.snp.makeConstraints { (make) in
            make.top.equalTo(view.safeAreaLayoutGuide.snp.top)
            make.width.equalTo(view.safeAreaLayoutGuide.snp.width)
            make.height.equalTo(view.safeAreaLayoutGuide.snp.height).multipliedBy(0.35)
            make.centerX.equalTo(view.safeAreaLayoutGuide.snp.centerX)
        }
        
        promptTextField.snp.makeConstraints { (make) in
            make.top.equalTo(generateView.snp.top).offset(5)
            make.width.equalTo(generateView.snp.width).multipliedBy(0.95)
            make.height.equalTo(generateView.snp.height).multipliedBy(0.35)
            make.centerX.equalTo(generateView.snp.centerX)
        }
        
        negativePromptTextField.snp.makeConstraints { (make) in
            make.top.equalTo(promptTextField.snp.bottom).offset(5)
            make.width.equalTo(generateView.snp.width).multipliedBy(0.95)
            make.height.equalTo(generateView.snp.height).multipliedBy(0.35)
            make.centerX.equalTo(generateView.snp.centerX)
        }
        
        generateButton.snp.makeConstraints { (make) in
            make.top.equalTo(negativePromptTextField.snp.bottom).offset(5)
            make.width.equalTo(generateView.snp.width).multipliedBy(0.95)
            make.height.equalTo(generateView.snp.height).multipliedBy(0.2)
            make.centerX.equalTo(generateView.snp.centerX)
        }
        
        generateImage.snp.makeConstraints { (make) in
            make.top.equalTo(generateView.snp.bottom).offset(10)
            make.width.equalTo(view.safeAreaLayoutGuide.snp.width)
            make.height.equalTo(view.safeAreaLayoutGuide.snp.height).multipliedBy(0.60)
            make.centerX.equalTo(view.safeAreaLayoutGuide.snp.centerX)
        }
    }
    
    func translatePrompt(prompt: String, negativePrompt: String) -> (tranPrompt: String, tranNegativePrompt: String){
        let APP_KEY = "3bb96f14d2af67ec"
        let APP_SECRET = "SRQF1jGwAMuYciFK33KfyRGKa7B1x38u"
        let q = prompt + "|" + negativePrompt
        let salt = UUID().uuidString
        let curtime = Int(Date().timeIntervalSince1970)
        print(curtime)
        let signStr = APP_KEY + q + salt + String(curtime) + APP_SECRET
        let cryptoSign = signStr.sha256
        let paramDic = ["q": q, "from": "auto", "to": "en", "appKey": APP_KEY, "salt": salt, "sign": cryptoSign, "signType": "v3", "curtime": String(curtime)]
        let list  = NSMutableArray()
        for subDic in paramDic {
            let tmpStr = "\(subDic.0)=\(subDic.1)"
            list.add(tmpStr)
        }
        let paramStr = list.componentsJoined(by: "&")
        let paramData = "https://openapi.youdao.com/api?" + paramStr
        var addr = paramData.addingPercentEncoding(withAllowedCharacters: .urlQueryAllowed)
        let url = URL(string: addr!)
        var request = URLRequest(url: url!)
        request.httpMethod = "GET"
        
        var promptValue = ""
        var negativePromotValue = ""
        
        let task = URLSession.shared.dataTask(with: request) { (data, response, error) in
            do {
                let jsonData = try JSON(data: data!)
                let promptData = jsonData["translation"][0].stringValue
                print(promptData)
                if String(promptData.first!) == "|" {
                    let subStringArr = promptData.split(separator: "| ")
                    negativePromotValue = String(subStringArr[0])
                    
                } else if String(promptData.last!) == "|" {
                    let subStringArr = promptData.split(separator: " |")
                    promptValue = String(subStringArr[0])
                } else {
                    let subStringArr = promptData.split(separator: " | ")
                    promptValue = String(subStringArr[0])
                    negativePromotValue = String(subStringArr[1])
                }
            } catch {}
        }
        task.resume()
        return (promptValue, negativePromotValue)
    }
    
    @objc func generateAction() {
        let translation = translatePrompt(prompt: promptTextField.text!, negativePrompt: negativePromptTextField.text!)
        generateButton.backgroundColor = .gray
        generateButton.isUserInteractionEnabled = false
        
        print("translation", translation)
        var promptValue = translation.tranPrompt
        var negativePromotValue = translation.tranNegativePrompt
        
        print("Generation: \(promptValue) + \(negativePromotValue)")
        
        prompt = "masterpiece, best quality, " + promptValue
        negativePrompt = "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digits, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, bad feet, " + negativePromotValue
        sessionHash = UUID().uuidString
        painting?.prompt = promptTextField.text!
        painting?.negative_prompt = negativePromptTextField.text!
        painting?.session_hash = sessionHash
        painting?.create_time = Date()
        let url = URL(string: "https://aipaint.hefengyu.org/api/predict/")
        var request = URLRequest(url: url!)
        request.httpMethod = "POST"
        let postData: [String: Any] = ["fn_index":50,"data":[prompt,negativePrompt,"None","None",20,"Euler a",false,false,1,1,7,-1,-1,0,0,0,false,512,512,false,0.7,0,0,"None",false,false,false,"","Seed","","Nothing","",true,false,false,"null","",""],"session_hash":sessionHash]
        let jsonData = try? JSONSerialization.data(withJSONObject: postData, options: [])
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        request.httpBody = jsonData
        var imageData: Data?
        let dataTask = URLSession.shared.dataTask(with: request) { (data, request, error) in
            do {
                let json = try? JSON(data: data!)
                let imageAddr = json!["data"][0][0]["name"].stringValue
                let addrProcess = imageAddr.replacingOccurrences(of: "\\", with: "/")
                let imageStr = "https://aipaint.hefengyu.org/file=\(addrProcess)"
                let imageUrl = URL(string: imageStr)
                imageData = try? Data(contentsOf: imageUrl!)
                DispatchQueue.main.async { [self] in
                    painting?.image = imageData!
                    self.generateImage.image = UIImage(data: imageData!)
                    self.generateButton.backgroundColor = .blue
                    self.generateButton.isUserInteractionEnabled = true
                    if context!.hasChanges {
                        do {
                            try context?.save()
                            print("save success.")
                        } catch {}
                    }
                }
            }
        }
//        let processURL = URL(string: "https://aipaint.hefengyu.org/api/predict/")
//        let processTask = URLSession.shared.dataTask(with: <#T##URLRequest#>)
        dataTask.resume()
    }
    
    @objc func saveImageToPhotos() {
        if (longPress!.state == .began) {
            // do something
            print("nice")
            guard let img = generateImage.image else {
                return
            }
            
            switch PHPhotoLibrary.authorizationStatus() {
            case .authorized:
                saveImg(image: img)
            case .notDetermined:
                PHPhotoLibrary.requestAuthorization { (status) in
                    if status == .authorized {
                        self.saveImg(image: img)
                    } else {
                        print("用户拒绝访问")
                    }
                }
            case .restricted, .denied:
                if let url = URL.init(string: UIApplication.openSettingsURLString) {
                    if UIApplication.shared.canOpenURL(url) {
                        UIApplication.shared.open(url)
                    }
                }
            case .limited:
                return
            @unknown default:
                return
            }
        }else if (longPress!.state == .ended){
            return
        }
    }
    
    func saveImg(image: UIImage) {
        PHPhotoLibrary.shared().performChanges({
            PHAssetChangeRequest.creationRequestForAsset(from: image)
        }, completionHandler: { (isSuccess, error) in
            DispatchQueue.main.async {
                if isSuccess{
                    self.present(self.alertController, animated: true, completion: nil)
                    DispatchQueue.main.asyncAfter(deadline: DispatchTime.now() + 2) {
                        self.alertController.dismiss(animated: false)
                    }
                }
            }
        })
    }
    
    override func touchesBegan(_ touches: Set<UITouch>, with event: UIEvent?) {
        promptTextField.resignFirstResponder()
        negativePromptTextField.resignFirstResponder()
    }
}