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

import Foundation

public final class AutoSongProjectDecisionsDirector: Sendable {
    private let aiService: any AIService

    public init(aiService: any AIService) {
        self.aiService = aiService
    }

    public func applyAutomaticDecisions(
        to project: SongProject,
        instruction: String? = nil,
        conversation: [AIConversationMessage] = [],
        localeIdentifier: String? = nil,
        updatedAt: Date = Date()
    ) async throws -> SongProjectUpdateResult {
        let allowedScopes = automaticScopes(for: project)
        guard !allowedScopes.isEmpty else {
            return SongProjectUpdateResult(
                project: project,
                notes: ["No automatic decisions are enabled."]
            )
        }

        let trimmedInstruction = instruction?.trimmingCharacters(in: .whitespacesAndNewlines)
        let request = SongProjectUpdateRequest(
            context: AIRequestContext(
                userInstruction: trimmedInstruction?.isEmpty == false
                    ? trimmedInstruction!
                    : "Complete the enabled automatic song-project decisions.",
                conversation: conversation,
                localeIdentifier: localeIdentifier
            ),
            project: project,
            allowedScopes: allowedScopes
        )
        let result = try await aiService.proposeProjectUpdate(from: request)
        var updatedProject = mergeAutomaticDecisions(
            from: result.project,
            into: project
        )
        updatedProject.updatedAt = updatedAt

        return SongProjectUpdateResult(
            project: updatedProject,
            changedScopes: changedScopes(from: project, to: updatedProject),
            notes: result.notes,
            followUpQuestions: result.followUpQuestions
        )
    }

    private func automaticScopes(for project: SongProject) -> [SongProjectUpdateScope] {
        var scopes: [SongProjectUpdateScope] = []

        if project.isStructureAutomatic {
            scopes.append(.structure)
        }

        if project.instruments.contains(where: { $0.selected && $0.autoArrangementEnabled }) {
            scopes.append(.arrangement)
        }

        if project.bpm?.mode != .manual ||
            project.key?.mode != .manual ||
            project.scale?.mode != .manual ||
            project.maqam?.mode != .manual {
            scopes.append(.musicalParameters)
        }

        if project.isProductionAutomatic {
            scopes.append(.productionDirections)
        }

        return scopes
    }

    private func mergeAutomaticDecisions(
        from proposedProject: SongProject,
        into project: SongProject
    ) -> SongProject {
        var updatedProject = project

        if project.isStructureAutomatic {
            updatedProject.sections = proposedProject.sections
        }

        updatedProject.instruments = mergedInstruments(
            current: project.instruments,
            proposed: proposedProject.instruments
        )
        updatedProject.bpm = automaticValue(current: project.bpm, proposed: proposedProject.bpm)
        updatedProject.key = automaticValue(current: project.key, proposed: proposedProject.key)
        updatedProject.scale = automaticValue(current: project.scale, proposed: proposedProject.scale)
        updatedProject.maqam = automaticValue(current: project.maqam, proposed: proposedProject.maqam)

        if project.isProductionAutomatic {
            updatedProject.productionDirections = proposedProject.productionDirections
            updatedProject.sections = sections(
                updatedProject.sections,
                applyingProductionDirectionsFrom: proposedProject.sections
            )
        } else {
            updatedProject.sections = sections(
                updatedProject.sections,
                preservingProductionDirectionsFrom: project.sections
            )
        }

        return updatedProject
    }

    private func mergedInstruments(
        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(
        _ current: [SongSection],
        preservingProductionDirectionsFrom existing: [SongSection]
    ) -> [SongSection] {
        var existingByID: [String: SongSection] = [:]
        for existingSection in existing {
            existingByID[existingSection.id] = existingSection
        }

        return current.map { section in
            var updatedSection = section
            updatedSection.productionDirection = existingByID[section.id]?.productionDirection
            return updatedSection
        }
    }

    private func changedScopes(
        from project: SongProject,
        to updatedProject: SongProject
    ) -> [SongProjectUpdateScope] {
        var scopes: [SongProjectUpdateScope] = []

        if project.isStructureAutomatic &&
            sectionsDifferIgnoringProductionDirections(project.sections, updatedProject.sections) {
            scopes.append(.structure)
        }
        if project.instruments != updatedProject.instruments {
            scopes.append(.arrangement)
        }
        if project.bpm != updatedProject.bpm ||
            project.key != updatedProject.key ||
            project.scale != updatedProject.scale ||
            project.maqam != updatedProject.maqam {
            scopes.append(.musicalParameters)
        }
        if project.isProductionAutomatic &&
            (project.productionDirections != updatedProject.productionDirections ||
                productionDirectionsDiffer(in: project.sections, and: updatedProject.sections)) {
            scopes.append(.productionDirections)
        }

        return scopes
    }

    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
        }
    }
}