photochk/Sources/Photochk/Core/SimilarityGrouper.swift

import Foundation
import Vision

/// 走査からグループ生成までの一連の処理をまとめる。
///
/// 安いフィルタから順にかけて候補を絞る:
/// 1. ファイルサイズが同じものだけ SHA-256 → 完全同一を確定
/// 2. dHash（知覚ハッシュ）で「見た目が近そうな」ペアを抽出
/// 3. そのペアに登場する画像だけ Vision の feature print で精査
/// 4. Union-Find で連結成分をグループ化
enum SimilarityGrouper {
    private struct Pair: Sendable {
        let a: Int
        let b: Int
    }

    private struct HashEntry: Sendable {
        let index: Int
        let hash: UInt64
    }

    /// feature print を計算できなかったときの保険。dHash がここまで一致していれば同一とみなす。
    private static let fallbackHammingThreshold = 4

    static func run(
        root: URL,
        settings: ScanSettings,
        cache: FingerprintCache,
        progress: @escaping @Sendable (ScanProgress) -> Void
    ) async throws -> [SimilarityGroup] {
        progress(ScanProgress(phase: .scanning))
        let files = try PhotoScanner.scan(root: root, includesSubfolders: settings.includesSubfolders)
        try Task.checkCancellation()
        guard files.count > 1 else { return [] }

        // --- Stage 1: 完全同一をまとめ、以降は各セットの代表1枚だけを扱う ---
        let leaderOf = try await exactLeaders(files: files, progress: progress)
        let leaderIndices = files.indices.filter { leaderOf[$0] == $0 }
        let leaderFiles = leaderIndices.map { files[$0] }

        // --- Stage 2: dHash ---
        let basicsByURL = try await loadBasics(for: leaderFiles, cache: cache, progress: progress)

        // --- Stage 3: 候補ペア → feature print ---
        let entries = leaderIndices.compactMap { index -> HashEntry? in
            guard let basics = basicsByURL[files[index].url] else { return nil }
            return HashEntry(index: index, hash: basics.dHash)
        }
        progress(ScanProgress(phase: .comparing))
        let pairs = try await candidatePairs(
            entries,
            threshold: settings.sensitivity.hammingThreshold
        )

        var involved = Set<Int>()
        for pair in pairs {
            involved.insert(pair.a)
            involved.insert(pair.b)
        }
        let involvedFiles = involved.sorted().map { files[$0] }
        let printsByURL = try await loadFeaturePrints(
            for: involvedFiles, cache: cache, progress: progress
        )

        // --- Stage 4: グループ化 ---
        progress(ScanProgress(phase: .grouping))
        var unionFind = UnionFind(count: files.count)
        for index in files.indices where leaderOf[index] != index {
            unionFind.union(index, leaderOf[index])
        }

        let threshold = settings.sensitivity.featurePrintThreshold
        for pair in pairs {
            let printA = printsByURL[files[pair.a].url]
            let printB = printsByURL[files[pair.b].url]

            if let distance = FeaturePrinter.distance(printA, printB) {
                if distance < threshold { unionFind.union(pair.a, pair.b) }
            } else if let basicsA = basicsByURL[files[pair.a].url],
                      let basicsB = basicsByURL[files[pair.b].url],
                      DHash.hammingDistance(basicsA.dHash, basicsB.dHash) <= fallbackHammingThreshold {
                unionFind.union(pair.a, pair.b)
            }
        }
        try Task.checkCancellation()

        let groups = buildGroups(
            components: unionFind.components(),
            files: files,
            leaderOf: leaderOf,
            basicsByURL: basicsByURL,
            printsByURL: printsByURL
        )

        await cache.prune(keeping: files)
        await cache.save()
        return groups
    }

    // MARK: - Stage 1

    /// 各インデックスに「完全同一セットの代表インデックス」を割り当てる。単独なら自分自身。
    private static func exactLeaders(
        files: [ScannedFile],
        progress: @escaping @Sendable (ScanProgress) -> Void
    ) async throws -> [Int] {
        var leaderOf = Array(files.indices)

        // サイズが違えば中身も違う。同サイズのものだけハッシュすればよい。
        var bySize: [Int64: [Int]] = [:]
        for (index, file) in files.enumerated() {
            bySize[file.fileSize, default: []].append(index)
        }
        let candidates = bySize.values.filter { $0.count > 1 }.flatMap { $0 }.sorted()
        guard !candidates.isEmpty else { return leaderOf }

        progress(ScanProgress(phase: .hashing, completed: 0, total: candidates.count))
        let hashes = try await mapConcurrently(
            candidates,
            onProgress: { done in
                progress(ScanProgress(phase: .hashing, completed: done, total: candidates.count))
            }
        ) { index -> String? in
            try? ContentHasher.sha256(of: files[index].url)
        }

        var byDigest: [String: [Int]] = [:]
        for (index, hash) in zip(candidates, hashes) {
            guard let hash else { continue }
            byDigest["\(files[index].fileSize):\(hash)", default: []].append(index)
        }
        for indices in byDigest.values where indices.count > 1 {
            guard let leader = indices.min(by: { files[$0].url.path < files[$1].url.path }) else {
                continue
            }
            for index in indices { leaderOf[index] = leader }
        }
        return leaderOf
    }

    // MARK: - Stage 2 / 3

    private static func loadBasics(
        for files: [ScannedFile],
        cache: FingerprintCache,
        progress: @escaping @Sendable (ScanProgress) -> Void
    ) async throws -> [URL: ImageBasics] {
        var result = await cache.basics(for: files)
        let missing = files.filter { result[$0.url] == nil }
        guard !missing.isEmpty else { return result }

        let cached = files.count - missing.count
        progress(ScanProgress(phase: .fingerprinting, completed: cached, total: files.count))
        let computed = try await mapConcurrently(
            missing,
            onProgress: { done in
                progress(
                    ScanProgress(phase: .fingerprinting, completed: cached + done, total: files.count)
                )
            }
        ) { file -> ImageBasics? in
            DHash.compute(for: file.url)
        }

        var fresh: [URL: ImageBasics] = [:]
        for (file, basics) in zip(missing, computed) {
            guard let basics else { continue }
            fresh[file.url] = basics
            result[file.url] = basics
        }
        await cache.store(basics: fresh, for: missing)
        return result
    }

    private static func loadFeaturePrints(
        for files: [ScannedFile],
        cache: FingerprintCache,
        progress: @escaping @Sendable (ScanProgress) -> Void
    ) async throws -> [URL: FeaturePrintObservation] {
        guard !files.isEmpty else { return [:] }
        var result = await cache.featurePrints(for: files)
        let missing = files.filter { result[$0.url] == nil }
        guard !missing.isEmpty else { return result }

        let cached = files.count - missing.count
        progress(ScanProgress(phase: .comparing, completed: cached, total: files.count))
        let computed = try await mapConcurrently(
            missing,
            onProgress: { done in
                progress(
                    ScanProgress(phase: .comparing, completed: cached + done, total: files.count)
                )
            }
        ) { file -> FeaturePrintObservation? in
            try? await FeaturePrinter.compute(for: file.url)
        }

        var fresh: [URL: FeaturePrintObservation] = [:]
        for (file, print) in zip(missing, computed) {
            guard let print else { continue }
            fresh[file.url] = print
            result[file.url] = print
        }
        await cache.store(featurePrints: fresh, for: missing)
        return result
    }

    /// dHash のハミング距離がしきい値以下のペアを総当たりで探す。
    ///
    /// 64bit の XOR + popcount なので1ペアあたりの実費はごく小さい。LSH などで
    /// 絞ると取りこぼしが出るため、区間をコアに分けて総当たりする方針にしている。
    private static func candidatePairs(
        _ entries: [HashEntry],
        threshold: Int
    ) async throws -> [Pair] {
        guard entries.count > 1 else { return [] }

        let cores = max(2, ProcessInfo.processInfo.activeProcessorCount)
        let chunkSize = max(64, entries.count / (cores * 4))
        let starts = stride(from: 0, to: entries.count, by: chunkSize).map { $0 }

        let chunks = try await mapConcurrently(starts) { start -> [Pair] in
            var found: [Pair] = []
            let end = min(start + chunkSize, entries.count)
            for i in start..<end {
                let hash = entries[i].hash
                for j in (i + 1)..<entries.count
                where DHash.hammingDistance(hash, entries[j].hash) <= threshold {
                    found.append(Pair(a: entries[i].index, b: entries[j].index))
                }
            }
            return found
        }
        return chunks.flatMap { $0 }
    }

    // MARK: - Stage 4

    /// グループの中に「代表とは別だが、互いにバイト単位で同じ」ファイルの組があれば印を付ける。
    ///
    /// 代表と同一のものは相手が「完全同一」と表示されるので印は要らない。ここで拾うのは
    /// たとえば縮小版が2つあってその2つだけが同一、といったケース。
    private static func exactSetTags(
        in ranked: [Int],
        leaderOf: [Int],
        excluding representativeLeader: Int
    ) -> [Int: String] {
        var membersPerLeader: [Int: Int] = [:]
        for index in ranked where leaderOf[index] != representativeLeader {
            membersPerLeader[leaderOf[index], default: 0] += 1
        }

        let labels = Array("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
        var tags: [Int: String] = [:]
        // ranked の順に印を振るので、同じ結果からは常に同じ印になる。
        for index in ranked {
            let leader = leaderOf[index]
            guard membersPerLeader[leader, default: 0] > 1, tags[leader] == nil else { continue }
            tags[leader] = tags.count < labels.count
                ? String(labels[tags.count])
                : "\(tags.count + 1)"
        }
        return tags
    }

    private static func buildGroups(
        components: [[Int]],
        files: [ScannedFile],
        leaderOf: [Int],
        basicsByURL: [URL: ImageBasics],
        printsByURL: [URL: FeaturePrintObservation]
    ) -> [SimilarityGroup] {
        func makeAsset(_ index: Int) -> PhotoAsset {
            let file = files[index]
            // 完全同一なら中身は同じなので、代表の解像度・撮影日をそのまま使える。
            let basics = basicsByURL[files[leaderOf[index]].url]
            return PhotoAsset(
                url: file.url,
                fileSize: file.fileSize,
                modificationDate: file.modificationDate,
                captureDate: basics?.captureDate ?? file.creationDate,
                pixelWidth: basics?.pixelWidth ?? 0,
                pixelHeight: basics?.pixelHeight ?? 0
            )
        }

        var groups: [SimilarityGroup] = []
        for component in components where component.count > 1 {
            let assets = Dictionary(uniqueKeysWithValues: component.map { ($0, makeAsset($0)) })

            // 代表は「解像度が最大 → ファイルサイズが最大 → 撮影が最古」で選ぶ。
            let ranked = component.sorted { lhs, rhs in
                guard let a = assets[lhs], let b = assets[rhs] else { return lhs < rhs }
                if a.pixelCount != b.pixelCount { return a.pixelCount > b.pixelCount }
                if a.fileSize != b.fileSize { return a.fileSize > b.fileSize }
                let dateA = a.captureDate ?? .distantFuture
                let dateB = b.captureDate ?? .distantFuture
                if dateA != dateB { return dateA < dateB }
                return a.url.path < b.url.path
            }
            guard let representativeIndex = ranked.first,
                  let representative = assets[representativeIndex]
            else { continue }
            let representativeLeader = leaderOf[representativeIndex]
            let exactSetTags = exactSetTags(
                in: ranked, leaderOf: leaderOf, excluding: representativeLeader
            )

            let members = component.compactMap { index -> SimilarityGroup.Member? in
                guard let asset = assets[index] else { return nil }
                let relation: SimilarityGroup.Relation
                if leaderOf[index] == representativeLeader {
                    relation = .exact
                } else {
                    relation = .similar(
                        distance: FeaturePrinter.distance(
                            printsByURL[files[leaderOf[index]].url],
                            printsByURL[files[representativeLeader].url]
                        )
                    )
                }
                return SimilarityGroup.Member(
                    asset: asset,
                    relation: relation,
                    exactSetTag: exactSetTags[leaderOf[index]]
                )
            }

            let ordered = members.sorted { lhs, rhs in
                if lhs.id == representative.url { return true }
                if rhs.id == representative.url { return false }
                if lhs.relation.sortKey != rhs.relation.sortKey {
                    return lhs.relation.sortKey < rhs.relation.sortKey
                }
                return lhs.asset.url.path < rhs.asset.url.path
            }

            groups.append(
                SimilarityGroup(members: ordered, representativeID: representative.url)
            )
        }

        // 減らせる容量が大きいグループから見せる。
        return groups.sorted { lhs, rhs in
            if lhs.reclaimableBytes != rhs.reclaimableBytes {
                return lhs.reclaimableBytes > rhs.reclaimableBytes
            }
            return lhs.representativeID.path < rhs.representativeID.path
        }
    }
}