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
import Alamofire

class DrawViewController: UIViewController, UIScrollViewDelegate, UIImagePickerControllerDelegate, UINavigationControllerDelegate {
    
    var drawModeView: UIScrollView!
    
    lazy var generateView: UIView = {
        let view = UIView()
        view.backgroundColor = UIColor(red: 48/255, green: 155/255, blue: 255/255, alpha: 1.0)
        view.layer.cornerRadius = 14.0
        return view
    }()
    
    lazy var promptTextField: UITextField = {
        let textField = UITextField()
        textField.placeholder = "请输入你的作品的组成元素，以空格分割。"
        textField.backgroundColor = .white
        textField.layer.cornerRadius = 14.0
        textField.returnKeyType = .done
        return textField
    }()
    
    lazy var negativePromptTextField: UITextField = {
        let textField = UITextField()
        textField.placeholder = "请输入你不想在作品中看到的元素，以空格分割。"
        textField.backgroundColor = .white
        textField.layer.cornerRadius = 14.0
        textField.returnKeyType = .done
        return textField
    }()
    
    lazy var generateButton: UIButton = {
        let button = UIButton()
        button.layer.backgroundColor = UIColor.black.cgColor
        button.addTarget(self, action: #selector(generateAction), for: .touchUpInside)
        button.setTitle("生成图片", for: .normal)
        button.tintColor = .white
        button.isUserInteractionEnabled = true
        button.layer.cornerRadius = 14.0
        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
    }()
    
    lazy var remindText: UILabel = {
        var label = UILabel()
        label.text = "长按图片进行保存"
        label.textColor = .systemGray5
        label.textAlignment = .center
        return label
    }()
    
    lazy var imageUploadView: UIView = {
        var view = UIView()
        view.layer.borderWidth = 1
        view.layer.borderColor = UIColor.black.cgColor
        return view
    }()
    
    lazy var pic2PicLabel: UILabel = {
        var label = UILabel()
        label.text = "图片动漫化"
        return label
    }()
    
    lazy var uploadImage: UIImageView = {
        var imageView = UIImageView()
        let singleGusture = UITapGestureRecognizer(target: self, action: #selector(pickPhotoFromAlbum))
        imageView.addGestureRecognizer(singleGusture)
        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)
    let userDefaults = UserDefaults.standard
    var pickedImage: UIImage?
    
    override func viewDidLoad() {
        super.viewDidLoad()
        drawModeView = UIScrollView()
        drawModeView.delegate = self
        drawModeView.showsVerticalScrollIndicator = true
        drawModeView.isScrollEnabled = true
        self.view.addSubview(drawModeView)

        drawModeView.snp.makeConstraints { (make) in
            make.edges.equalTo(view.safeAreaLayoutGuide.snp.edges)
        }
    
        if UITraitCollection.current.userInterfaceStyle == .dark {
            view.backgroundColor = .black
            generateView.layer.borderColor = UIColor.white.cgColor
            promptTextField.layer.borderColor = UIColor.white.cgColor
            negativePromptTextField.layer.borderColor = UIColor.white.cgColor
            generateImage.layer.borderColor = UIColor.white.cgColor
            remindText.textColor = .white
        } else {
            view.backgroundColor = .white
            generateView.layer.borderColor = UIColor.black.cgColor
            promptTextField.layer.borderColor = UIColor.black.cgColor
            negativePromptTextField.layer.borderColor = UIColor.black.cgColor
            generateImage.layer.borderColor = UIColor.black.cgColor
            remindText.textColor = .systemGray5
        }
        self.title = "画酱"
        delegate = UIApplication.shared.delegate as? AppDelegate

        drawModeView.addSubview(generateView)
        generateView.addSubview(promptTextField)
        generateView.addSubview(negativePromptTextField)
        generateView.addSubview(generateButton)
        
        longPress = UILongPressGestureRecognizer(target: self, action: #selector(saveImageToPhotos))
        generateImage.addGestureRecognizer(longPress!)
        
        generateView.snp.makeConstraints { (make) in
            make.top.equalTo(drawModeView.contentLayoutGuide.snp.top)
            make.centerX.equalTo(view.safeAreaLayoutGuide.snp.centerX)
            make.left.equalTo(drawModeView.contentLayoutGuide.snp.left).offset(5)
            make.right.equalTo(drawModeView.contentLayoutGuide.snp.right).offset(-5)
            make.height.equalTo(drawModeView.contentLayoutGuide.snp.width).multipliedBy(0.75)
        }
        
        promptTextField.snp.makeConstraints { (make) in
            make.top.equalTo(generateView.snp.top).offset(5)
            make.left.equalTo(generateView.snp.left).offset(5)
            make.right.equalTo(generateView.snp.right).offset(-5)
            make.height.equalTo(generateView.snp.height).multipliedBy(0.4)
            make.centerX.equalTo(generateView.snp.centerX)
        }
        
        promptTextField.delegate = self
        
        negativePromptTextField.snp.makeConstraints { (make) in
            make.top.equalTo(promptTextField.snp.bottom).offset(5)
            make.left.equalTo(generateView.snp.left).offset(5)
            make.right.equalTo(generateView.snp.right).offset(-5)
            make.height.equalTo(generateView.snp.height).multipliedBy(0.32)
            make.centerX.equalTo(generateView.snp.centerX)
        }
        
        negativePromptTextField.delegate = self
        
        generateButton.snp.makeConstraints { (make) in
            make.top.equalTo(negativePromptTextField.snp.bottom).offset(5)
            make.width.equalTo(generateView.snp.width).multipliedBy(0.97)
            make.bottom.equalTo(generateView.snp.bottom).offset(-5)
            make.centerX.equalTo(generateView.snp.centerX)
        }

        drawModeView.addSubview(imageUploadView)
        imageUploadView.addSubview(pic2PicLabel)
        imageUploadView.addSubview(uploadImage)

        imageUploadView.snp.makeConstraints { (make) in
            make.top.equalTo(generateView.snp.bottom).offset(5)
            make.left.equalTo(drawModeView.contentLayoutGuide.snp.left).offset(5)
            make.right.equalTo(drawModeView.contentLayoutGuide.snp.right).offset(-5)
            make.centerX.equalTo(drawModeView.contentLayoutGuide.snp.centerX)
            make.height.equalTo(imageUploadView.snp.width).multipliedBy(0.4)
        }
        
        pic2PicLabel.snp.makeConstraints { (make) in
            make.left.equalTo(imageUploadView.snp.left).offset(5)
            make.top.equalTo(imageUploadView.snp.top).offset(5)
            make.width.equalTo(imageUploadView.snp.width)
            make.height.equalTo(25)
        }

        uploadImage.snp.makeConstraints { (make) in
            make.top.equalTo(pic2PicLabel.snp.bottom).offset(5)
            make.bottom.equalTo(imageUploadView.snp.bottom)
            make.left.equalTo(imageUploadView.snp.left).offset(5)
            make.right.equalTo(imageUploadView.snp.right).offset(-5)
        }
        
        drawModeView.addSubview(remindText)
        drawModeView.addSubview(generateImage)

        generateImage.snp.makeConstraints { (make) in
            make.top.equalTo(imageUploadView.snp.bottom).offset(10)
            make.left.equalTo(drawModeView.contentLayoutGuide.snp.left).offset(5)
            make.right.equalTo(drawModeView.contentLayoutGuide.snp.right).offset(-5)
            make.centerX.equalTo(drawModeView.contentLayoutGuide.snp.centerX)
            make.height.equalTo(generateImage.snp.width)
            make.bottom.equalTo(drawModeView.contentLayoutGuide.snp.bottom)
        }

        remindText.snp.makeConstraints { (make) in
            make.top.equalTo(generateImage.snp.top).offset(2)
            make.width.equalTo(generateImage.snp.width)
            make.height.equalTo(generateImage.snp.height).multipliedBy(0.40)
            make.centerX.equalTo(generateImage.snp.centerX)
        }
    }
    
    override func viewWillAppear(_ animated: Bool) {
        if !userDefaults.bool(forKey: UserDefaultKeys.UserInfo.isLogin) {
            generateButton.setTitle("请登录后使用", for: .normal)
            generateButton.isEnabled = false
        } else {
            generateButton.setTitle("生成图片", for: .normal)
            generateButton.isEnabled = true
        }
    }
    
    func translatePrompt(prompt: String, negativePrompt: String, completion: @escaping (_ tranPrompt: String, _ tranNegativePrompt: String) -> Void) {
        let APP_KEY = "3bb96f14d2af67ec"
        let APP_SECRET = "SRQF1jGwAMuYciFK33KfyRGKa7B1x38u"
        let q = prompt + "@" + negativePrompt
        let salt = UUID().uuidString
        let curtime = Int(Date().timeIntervalSince1970)
        let inputQ = q.count <= 20 ? q : String(q.prefix(10)) + String(q.count) + String(q.suffix(10))
        let signStr = APP_KEY + inputQ + 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 = APIs.YouDao.TextTranslate + "api?" + paramStr
        print(paramData)
        guard let addr = paramData.addingPercentEncoding(withAllowedCharacters: .urlQueryAllowed) else {
            print("error")
            return
        }
        guard let url = URL(string: addr) else {
            print("error")
            return
        }
        var request = URLRequest(url: url)
        request.httpMethod = "GET"
        
        var promptValue = ""
        var negativePromotValue = ""
        
        let task = URLSession.shared.dataTask(with: request) { (data, response, error) in
            do {
                guard let unwrapedData = data else {
                    print("error")
                    return
                }
                let jsonData = try JSON(data: unwrapedData)
                let promptData = jsonData["translation"][0].stringValue
                print(promptData)
                if String(promptData) == "@" {
                    
                } else {
                    if String(promptData.first ?? String.Element("")) == "@" {
                        if #available(iOS 16.0, *) {
                            let subStringArr = promptData.split(separator: "@")
                            negativePromotValue = String(subStringArr[0])
                        } else {
                            // Fallback on earlier versions
                            let subStringArr = promptData.components(separatedBy: "@")
                            negativePromotValue = String(subStringArr[0])
                        }
                        
                    } else if String(promptData.last ?? String.Element("")) == "@" {
                        if #available(iOS 16.0, *) {
                            let subStringArr = promptData.split(separator: " @")
                            promptValue = String(subStringArr[0])
                        } else {
                            // Fallback on earlier versions
                            let subStringArr = promptData.components(separatedBy: "@")
                            promptValue = String(subStringArr[0])
                        }
                    } else if promptData == "@"{
                        // do nothing
                    } else {
                        if #available(iOS 16.0, *) {
                            let subStringArr = promptData.split(separator: "@")
                            promptValue = String(subStringArr[0])
                            negativePromotValue = String(subStringArr[1])
                        } else {
                            // Fallback on earlier versions
                            let subStringArr = promptData.components(separatedBy: "@")
                            promptValue = String(subStringArr[0])
                            negativePromotValue = String(subStringArr[1])
                        }
                        
                    }
                }
            } catch {}
            completion(promptValue, negativePromotValue)
        }
        task.resume()
    }
    
    @objc func generateAction() {
        let promptText = promptTextField.text ?? ""
        let negativePromptText = negativePromptTextField.text ?? ""

        translatePrompt(prompt: promptTextField.text!, negativePrompt: negativePromptTextField.text!) { [self] promptValue, negativePromotValue in
            
            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.replacingOccurrences(of: "-", with: "").lowercased()
            guard let unwrapedDelegate = delegate else {
                return
            }
            context = unwrapedDelegate.persistentContainer.viewContext
            painting = NSEntityDescription.insertNewObject(forEntityName: "Painting", into: context!) as? Painting
            painting?.prompt = promptText
            painting?.negative_prompt = negativePromptText
            painting?.session_hash = sessionHash
            painting?.create_time = Date()

            guard let userID = userDefaults.string(forKey: UserDefaultKeys.UserInfo.userId) else {
                return
            }
            if pickedImage != nil {
                let imgBase64 = pickedImage?.pngData()?.base64EncodedString()
                
                let imgStr = "data:image/png;base64," + imgBase64!
                let postData: [String: String] = ["userId": userID, "img": imgStr, "prompt": prompt, "negPrompt": negativePrompt, "sessionHash": sessionHash]
                AF.request(APIs.Draw.Pic2Pic, method: .post, parameters: postData, encoder: URLEncodedFormParameterEncoder.default).response { [self] response in
                    switch response.result {
                    case .success(let data):
                        guard let json = try? JSON(data: data!) else {
                            return
                        }
                        let imageAddr = json["result"].stringValue
                        let addrProcess = imageAddr.replacingOccurrences(of: "\\", with: "/")
                        let url = URL(string: addrProcess)
                        generateImage.downloadedFrom(link: addrProcess)
                        let dataTask = URLSession.shared.dataTask(with: url!) { [self] (data, response, error) in
                            painting?.image = data
                            saveContext()
                        }
                        dataTask.resume()
                        DispatchQueue.main.async { [self] in
                                self.generateButton.backgroundColor = .blue
                                self.generateButton.isUserInteractionEnabled = true
                        }
                        break
                    case .failure(let error):
                        print(error)
                        break
                    }
                }
            } else {
                let postData: [String: String] = ["userId": userID, "prompt": prompt, "negPrompt": negativePrompt, "sessionHash": sessionHash]
                
                AF.request(APIs.Draw.Text2Pic, method: .post, parameters: postData, encoder: URLEncodedFormParameterEncoder.default).response { [self] response in
                    switch response.result {
                    case .success(let data):
                        guard let json = try? JSON(data: data!) else {
                            return
                        }
                        let imageAddr = json["result"].stringValue
                        let addrProcess = imageAddr.replacingOccurrences(of: "\\", with: "/")
                        let url = URL(string: addrProcess)
                        generateImage.downloadedFrom(link: addrProcess)
                        let dataTask = URLSession.shared.dataTask(with: url!) { [self] (data, response, error) in
                            painting?.image = data
                            saveContext()
                        }
                        dataTask.resume()
                        DispatchQueue.main.async { [self] in
                                self.generateButton.backgroundColor = .blue
                                self.generateButton.isUserInteractionEnabled = true
                        }
                        break
                    case .failure(let error):
                        print(error)
                        break
                    }
                }
            }
            
        }
        generateButton.backgroundColor = .gray
        generateButton.isUserInteractionEnabled = false
    }
    
    var allowEditing = true

    @objc func pickPhotoFromAlbum() {
        if UIImagePickerController.isSourceTypeAvailable(.photoLibrary) {
            let picker = UIImagePickerController()
            picker.delegate = self
            picker.sourceType = .photoLibrary
            picker.allowsEditing = allowEditing
            self.present(picker, animated: true)
        } else {
            print("读取相册时出错！")
        }
    }
    
    func imagePickerController(_ picker: UIImagePickerController, didFinishPickingMediaWithInfo info: [UIImagePickerController.InfoKey : Any]) {
        let image: UIImage!
        if allowEditing {
            image = info[.editedImage] as? UIImage
        } else {
            image = info[.originalImage] as? UIImage
        }
        uploadImage.image = image
        pickedImage = image
        picker.dismiss(animated: true)
    }
    
    @objc func saveImageToPhotos() {
        if (longPress!.state == .began) {
            guard let img = generateImage.image else {
                print("error")
                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 saveContext() {
        if context!.hasChanges {
            do {
                try context?.save()
                print("save success.")
            } catch {}
        }
    }
    
    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()
    }
    
    override func traitCollectionDidChange(_ previousTraitCollection: UITraitCollection?) {
        if #available(iOS 13, *) {
                if self.traitCollection.hasDifferentColorAppearance(comparedTo: previousTraitCollection) {
                    if self.traitCollection.userInterfaceStyle == .dark {
                        // 黑夜模式
                        view.backgroundColor = .black
                        generateView.layer.borderColor = UIColor.white.cgColor
                        promptTextField.layer.borderColor = UIColor.white.cgColor
                        negativePromptTextField.layer.borderColor = UIColor.white.cgColor
                        generateImage.layer.borderColor = UIColor.white.cgColor
                        remindText.textColor = .white
                    } else {
                        // 白天模式
                        view.backgroundColor = .white
                        generateView.layer.borderColor = UIColor.black.cgColor
                        promptTextField.layer.borderColor = UIColor.black.cgColor
                        negativePromptTextField.layer.borderColor = UIColor.black.cgColor
                        generateImage.layer.borderColor = UIColor.black.cgColor
                        remindText.textColor = .systemGray5
                    }
                }
            }
    }
}

extension DrawViewController: UITextFieldDelegate {
    func textFieldShouldReturn(_ textField: UITextField) -> Bool {
        self.view?.endEditing(false)
        return true
    }
}