music-assistant/Sources/MusicAssistantCore/Services/AI/SongProjectAIUpdateMerger.swift

import Foundation

public struct SongProjectAIUpdateMerger: Sendable {
    public init() {}

    public func apply(
        snapshot: SongProjectSnapshot,
        changedScopes: [SongProjectUpdateScope],
        notes: [String] = [],
        followUpQuestions: [String] = [],
        to existingProject: SongProject,
        updatedAt: Date = Date()
    ) -> SongProjectUpdateResult {
        var updatedProject = existingProject
        var appliedScopes: [SongProjectUpdateScope] = []

        for scope in uniqueScopes(changedScopes) where !existingProject.isUpdateScopeLocked(scope) {
            let projectBeforeScope = updatedProject
            apply(scope, snapshot: snapshot, to: &updatedProject)

            if didChange(scope, from: projectBeforeScope, to: updatedProject) {
                appliedScopes.append(scope)
            }
        }

        updatedProject.updatedAt = updatedAt
        return SongProjectUpdateResult(
            project: updatedProject,
            changedScopes: appliedScopes,
            notes: notes,
            followUpQuestions: followUpQuestions
        )
    }

    private func apply(
        _ scope: SongProjectUpdateScope,
        snapshot: SongProjectSnapshot,
        to project: inout SongProject
    ) {
        switch scope {
        case .duration:
            project.duration = snapshot.duration
        case .genres:
            project.genres = snapshot.genres
        case .instruments:
            project.instruments = mergedInstrumentSelections(
                current: project.instruments,
                proposed: snapshot.instruments
            )
        case .vocals:
            project.vocalists = snapshot.vocalists
        case .languages:
            project.languages = snapshot.languages
            project.dialects = snapshot.dialects
            project.arabicPronunciationSettings = snapshot.arabicPronunciationSettings
        case .lyrics:
            project.lyrics = snapshot.lyrics
        case .structure:
            guard project.isStructureAutomatic else { return }
            project.sections = sections(
                snapshot.sections,
                preservingProductionDirectionsFrom: project.sections
            )
        case .arrangement:
            project.instruments = mergedArrangement(
                current: project.instruments,
                proposed: snapshot.instruments
            )
        case .musicalParameters:
            project.bpm = automaticValue(current: project.bpm, proposed: snapshot.bpm)
            project.key = automaticValue(current: project.key, proposed: snapshot.key)
            project.scale = automaticValue(current: project.scale, proposed: snapshot.scale)
            project.maqam = automaticValue(current: project.maqam, proposed: snapshot.maqam)
        case .emotionalArc:
            project.emotionalArc = snapshot.emotionalArc
        case .productionDirections:
            guard project.isProductionAutomatic else { return }
            project.productionDirections = snapshot.productionDirections
            project.sections = sections(
                project.sections,
                applyingProductionDirectionsFrom: snapshot.sections
            )
        case .sunoOutput:
            project.sunoOutput = snapshot.sunoOutput
        }
    }

    private func mergedInstrumentSelections(
        current: [InstrumentTrack],
        proposed: [InstrumentTrack]
    ) -> [InstrumentTrack] {
        var proposedByID: [String: InstrumentTrack] = [:]
        for proposedTrack in proposed {
            proposedByID[proposedTrack.instrumentId] = proposedTrack
        }

        let mergedTracks = current.map { currentTrack in
            guard currentTrack.autoArrangementEnabled,
                  let proposedTrack = proposedByID[currentTrack.instrumentId]
            else {
                return currentTrack
            }

            var mergedTrack = currentTrack
            mergedTrack.selected = proposedTrack.selected
            mergedTrack.variant = proposedTrack.variant
            return mergedTrack
        }

        var knownInstrumentIDs = Set(current.map(\.instrumentId))
        let newTracks = proposed.filter { knownInstrumentIDs.insert($0.instrumentId).inserted }
        return mergedTracks + newTracks
    }

    private func mergedArrangement(
        current: [InstrumentTrack],
        proposed: [InstrumentTrack]
    ) -> [InstrumentTrack] {
        var proposedByID: [String: InstrumentTrack] = [:]
        for proposedTrack in proposed {
            proposedByID[proposedTrack.instrumentId] = proposedTrack
        }

        return current.map { currentTrack in
            guard currentTrack.selected,
                  currentTrack.autoArrangementEnabled,
                  let proposedTrack = proposedByID[currentTrack.instrumentId]
            else {
                return currentTrack
            }

            var mergedTrack = currentTrack
            mergedTrack.playingStyle = proposedTrack.playingStyle
            mergedTrack.role = proposedTrack.role
            mergedTrack.placements = proposedTrack.placements
            return mergedTrack
        }
    }

    private func automaticValue<Value>(
        current: ManualAutoValue<Value>?,
        proposed: ManualAutoValue<Value>?
    ) -> ManualAutoValue<Value>? where Value: Codable & Equatable & Sendable {
        guard current?.mode != .manual, let proposed else {
            return current
        }

        return ManualAutoValue(mode: .auto, value: proposed.value)
    }

    private func sections(
        _ current: [SongSection],
        applyingProductionDirectionsFrom proposed: [SongSection]
    ) -> [SongSection] {
        var proposedByID: [String: SongSection] = [:]
        for proposedSection in proposed {
            proposedByID[proposedSection.id] = proposedSection
        }

        return current.map { section in
            guard let proposedSection = proposedByID[section.id] else { return section }
            var updatedSection = section
            updatedSection.productionDirection = proposedSection.productionDirection
            return updatedSection
        }
    }

    private func sections(
        _ proposed: [SongSection],
        preservingProductionDirectionsFrom current: [SongSection]
    ) -> [SongSection] {
        var currentByID: [String: SongSection] = [:]
        for currentSection in current {
            currentByID[currentSection.id] = currentSection
        }

        return proposed.map { section in
            var updatedSection = section
            updatedSection.productionDirection = currentByID[section.id]?.productionDirection
            return updatedSection
        }
    }

    private func didChange(
        _ scope: SongProjectUpdateScope,
        from current: SongProject,
        to updated: SongProject
    ) -> Bool {
        switch scope {
        case .duration:
            return current.duration != updated.duration
        case .genres:
            return current.genres != updated.genres
        case .instruments, .arrangement:
            return current.instruments != updated.instruments
        case .vocals:
            return current.vocalists != updated.vocalists
        case .languages:
            return current.languages != updated.languages ||
                current.dialects != updated.dialects ||
                current.arabicPronunciationSettings != updated.arabicPronunciationSettings
        case .lyrics:
            return current.lyrics != updated.lyrics
        case .structure:
            return sectionsDifferIgnoringProductionDirections(current.sections, updated.sections)
        case .musicalParameters:
            return current.bpm != updated.bpm ||
                current.key != updated.key ||
                current.scale != updated.scale ||
                current.maqam != updated.maqam
        case .emotionalArc:
            return current.emotionalArc != updated.emotionalArc
        case .productionDirections:
            return current.productionDirections != updated.productionDirections ||
                productionDirectionsDiffer(in: current.sections, and: updated.sections)
        case .sunoOutput:
            return current.sunoOutput != updated.sunoOutput
        }
    }

    private func uniqueScopes(_ scopes: [SongProjectUpdateScope]) -> [SongProjectUpdateScope] {
        var seen = Set<SongProjectUpdateScope>()
        return scopes.filter { seen.insert($0).inserted }
    }

    private func sectionsDifferIgnoringProductionDirections(
        _ current: [SongSection],
        _ updated: [SongSection]
    ) -> Bool {
        guard current.count == updated.count else { return true }

        return zip(current, updated).contains { currentSection, updatedSection in
            var currentWithoutProduction = currentSection
            var updatedWithoutProduction = updatedSection
            currentWithoutProduction.productionDirection = nil
            updatedWithoutProduction.productionDirection = nil
            return currentWithoutProduction != updatedWithoutProduction
        }
    }

    private func productionDirectionsDiffer(
        in current: [SongSection],
        and updated: [SongSection]
    ) -> Bool {
        guard current.count == updated.count else { return true }

        return zip(current, updated).contains { currentSection, updatedSection in
            currentSection.productionDirection != updatedSection.productionDirection
        }
    }
}