clean_pro_photo_classifier_plugin/ios/Classes/ImageClassfiy/Models.swift

//
//  Models.swift
//  Demo
//
//  Created by Groot on 2025/4/27.
//

import UIKit
import Vision
import CoreLocation
import Photos

/// 分类类型选项集
public enum PhotoImageClassifyType: String, CaseIterable {
    case similar, people, screenshot, blurry

    public static var all: [PhotoImageClassifyType] {
        return Self.allCases
    }

    public init?(_ rawValue: String) {
        switch rawValue.lowercased() {
        case "similar":
            self = .similar
        case "people":
            self = .people
        case "screenshot":
            self = .screenshot
        case "blurry":
            self = .blurry
        default:
            return nil
        }
    }
}

/// 进度更新信息
public struct ClassificationProgress: Codable {
    /// 当前处理的批次索引
    public let currentBatch: Int
    /// 总批次数
    public let totalBatches: Int
    /// 进度百分比 (0.0-1.0)
    public let rate: Double
    /// 当前批次处理时间（秒）
    public let batchDuration: Double
    /// 总耗时（秒）
    public let totalDuration: Double
    /// 是否为最后一个批次
    public let isLastBatch: Bool
    
    public var isCompleted: Bool {
        return isLastBatch
    }
}

/// 回调数据
public struct PhotoImageClassifiedResult: Identifiable {
    public let id: String = UUID().uuidString
    public var similarGroups: [ImageGroup]?
    public var peopleImages: [ImageItem]?
    public var screenshotImages: [ImageItem]?
    public var blurryImages: [ImageItem]?
}

/// 图片组
public struct ImageGroup: Identifiable {
    public let id = UUID()
    
    public var images: [ImageItem]
    
    public var imageCount: Int {
        images.count
    }
    
    public var imageSizeCount: Int64 {
        images.reduce(0) { $0 + ($1.imageInfo?.fileSize ?? 0) }
    }
}

/// 图片项
public struct ImageItem: Identifiable {
    public let id = UUID()
    public let asset: PHAsset
    public var thumbnailImage: UIImage?
    
    public var imageInfo: ImageInfo?
    
    public var classfiyInfo: ClassifyInfo?
    
    public struct ImageInfo {
        public let time: Date?
        public let locationCoordinate: CLLocationCoordinate2D?
        public let fileSize: Int64?
        public let pixelSize: CGSize
    }
    
    public struct ClassifyInfo {
        public var similarFeature: VNFeaturePrintObservation? = nil
        
        public var containsPeople: Bool = false
        
        public var isBlurry: Bool = false
        
        public var isScreenshot: Bool = false
    }

    public func toDictionary() -> [String: Any] {
        var dict: [String: Any] = [
            "assetsId": asset.localIdentifier,
            "createdAt": imageInfo?.time?.timeIntervalSince1970 ?? 0,
            "fileSize": imageInfo?.fileSize ?? 0,
            "pixelWidth": Int(imageInfo?.pixelSize.width ?? 0),
            "pixelHeight": Int(imageInfo?.pixelSize.height ?? 0)
        ]
        
        if let coordinate = imageInfo?.locationCoordinate {
            dict["location"] = [
                "latitude": coordinate.latitude,
                "longitude": coordinate.longitude
            ]
        }
        return dict
    }
}