clean_pro_photo_classifier_plugin/ios/Classes/PhotoClassifierPlugin.swift

import Flutter
import UIKit
import Photos

public class PhotoClassifierPlugin: NSObject, FlutterPlugin, FlutterStreamHandler {
    private var eventSink: FlutterEventSink?
    
    public static func register(with registrar: FlutterPluginRegistrar) {
        // 注册方法通道
        let methodChannel = FlutterMethodChannel(name: "com.atmob.photo_classifier", binaryMessenger: registrar.messenger())
        // 注册事件通道
        let eventChannel = FlutterEventChannel(name: "com.atmob.photo_classifier_events", binaryMessenger: registrar.messenger())
        
        let instance = PhotoClassifierPlugin()
        registrar.addMethodCallDelegate(instance, channel: methodChannel)
        eventChannel.setStreamHandler(instance)
    }
    
    public func handle(_ call: FlutterMethodCall, result: @escaping FlutterResult) {
        // 确保我们在主线程上处理UI相关操作
        DispatchQueue.main.async {
            switch call.method {
            case "startClassification":
                self.handleStartClassification(call: call, result: result)
                
            case "configureClassifier":
                self.handleConfigureClassifier(call: call, result: result)
                
            case "resetClassifier":
                PhotosImageClassifier.shared.reset()
                result(nil)
                
            case "checkAssetsLoaded":
                result(PhotosImageClassifier.shared.isAssetsLoaded)
                
            default:
                result(FlutterMethodNotImplemented)
            }
        }
    }
}

// MARK: -- Platform Methods
extension PhotoClassifierPlugin {
    // 处理开始分类请求
    private func handleStartClassification(call: FlutterMethodCall, result: @escaping FlutterResult) {
        guard let args = call.arguments as? [String: Any], let types = self.convertClassifyTypes(args["types"] as? [String]), !types.isEmpty else {
            result(FlutterError(code: "INVALID_ARGUMENTS",
                                message: "Missing or invalid types argument",
                                details: nil))
            result(nil)
            return
        }
        
        Task {
            // 开始分类过程
            await PhotosImageClassifier.shared.startClassification(with: types) { [weak self] progress, classifiedResult in
                guard let self = self, let eventSink = self.eventSink else { return }
                
                // 转换进度信息
                let progressDict = progress.toDictionary()
                
                // 转换分类结果
                let resultDict = self.convertClassifiedResultToDictionary(classifiedResult)
                
                // 发送进度更新
                let updateDict: [String: Any] = [
                    "progress": progressDict,
                    "result": resultDict
                ]
                await MainActor.run {
                    eventSink(updateDict)
                    if progress.isCompleted {
                        eventSink(FlutterEndOfEventStream)
                    }
                }
            }
            result(nil)
        }
    }
    
    // 处理配置分类器
    private func handleConfigureClassifier(call: FlutterMethodCall, result: @escaping FlutterResult) {
        defer { result(nil) }
        guard let args = call.arguments as? [String: Any] else {
            result(FlutterError(code: "INVALID_ARGUMENTS",
                                message: "Missing arguments",
                                details: nil))
            return
        }
        // 解码参数
        guard let configData = try? JSONSerialization.data(withJSONObject: args),
              let config = try? JSONDecoder().decode(PhotosImageClassifier.Configuration.self, from: configData) else {
            result(FlutterError(code: "INVALID_CONFIG",
                              message: "无法解析配置参数",
                              details: nil))
            return
        }
        // 应用配置
        PhotosImageClassifier.shared.configure(with: config)
    }
    
    // 将字符串数组转换为分类类型
    private func convertClassifyTypes(_ types: [String]?) -> [PhotoImageClassifyType]? {
        guard let types else { return nil }
        var classifyTypes: Set<PhotoImageClassifyType> = []
        
        for type in types {
            guard let classifyType = PhotoImageClassifyType(type) else { continue }
            classifyTypes.insert(classifyType)
        }
        return Array(classifyTypes)
    }
    
    // 将分类结果转换为字典
    private func convertClassifiedResultToDictionary(_ result: PhotoImageClassifiedResult) -> [String: Any] {
        var resultDict: [String: Any] = [:]
        
        // 处理相似组
        if let similarGroups = result.similarGroups {
            resultDict["similarGroups"] = [
                "groups": similarGroups.map { group in
                     group.images.map { $0.toDictionary() }
                }
            ]
        }
        
        // 处理人物图片
        if let peopleImages = result.peopleImages {
            resultDict["peopleImages"] = peopleImages.map { $0.toDictionary() }
        }
        
        // 处理截图
        if let screenshotImages = result.screenshotImages {
            resultDict["screenshotImages"] = screenshotImages.map { $0.toDictionary() }
        }
        
        // 处理模糊图片
        if let blurryImages = result.blurryImages {
            resultDict["blurryImages"] = blurryImages.map { $0.toDictionary() }
        }
        
        return resultDict
    }
    
}

extension PhotoClassifierPlugin {
    public func onListen(withArguments arguments: Any?, eventSink events: @escaping FlutterEventSink) -> FlutterError? {
        self.eventSink = events
        return nil
    }
    
    public func onCancel(withArguments arguments: Any?) -> FlutterError? {
        self.eventSink = nil
        return nil
    }
}