music-assistant/Sources/MusicAssistantCore/Services/SongGeneration/SongGenerationContext.swift

import Foundation

public struct SongGenerationContext: Codable, Equatable, Sendable {
    public static let empty = SongGenerationContext(selectedInstruments: [])

    public var selectedInstruments: [SelectedInstrumentContext]

    public init(selectedInstruments: [SelectedInstrumentContext]) {
        self.selectedInstruments = selectedInstruments
    }

    public init(project: SongProject) {
        let sectionsByID = Dictionary(uniqueKeysWithValues: project.sections.map { ($0.id, $0) })
        selectedInstruments = project.selectedInstrumentTracks.map { track in
            SelectedInstrumentContext(
                instrumentId: track.instrumentId,
                displayName: Self.displayName(for: track),
                playingStyle: Self.normalized(track.playingStyle),
                role: Self.normalized(track.role),
                autoArrangementEnabled: track.autoArrangementEnabled,
                placements: track.placements.map { placement in
                    SelectedInstrumentPlacementContext(
                        sectionId: placement.sectionId,
                        sectionTitle: placement.sectionId.flatMap { sectionsByID[$0]?.title },
                        sectionType: placement.sectionId.flatMap { sectionsByID[$0]?.type },
                        startTime: placement.startTime,
                        endTime: placement.endTime,
                        direction: Self.normalized(placement.direction)
                    )
                }
            )
        }
    }

    public var hasSelectedInstruments: Bool {
        !selectedInstruments.isEmpty
    }

    public var sunoStyleInstrumentPhrase: String {
        selectedInstruments
            .map(\.stylePromptPhrase)
            .joined(separator: ", ")
    }

    private static func displayName(for track: InstrumentTrack) -> String {
        normalized(track.variant) ?? track.instrumentId
    }

    private static func normalized(_ value: String?) -> String? {
        let trimmedValue = value?.trimmingCharacters(in: .whitespacesAndNewlines)
        return trimmedValue?.isEmpty == false ? trimmedValue : nil
    }
}

public struct SelectedInstrumentContext: Codable, Equatable, Sendable {
    public var instrumentId: String
    public var displayName: String
    public var playingStyle: String?
    public var role: String?
    public var autoArrangementEnabled: Bool
    public var placements: [SelectedInstrumentPlacementContext]

    public init(
        instrumentId: String,
        displayName: String,
        playingStyle: String? = nil,
        role: String? = nil,
        autoArrangementEnabled: Bool = true,
        placements: [SelectedInstrumentPlacementContext] = []
    ) {
        self.instrumentId = instrumentId
        self.displayName = displayName
        self.playingStyle = playingStyle
        self.role = role
        self.autoArrangementEnabled = autoArrangementEnabled
        self.placements = placements
    }

    public var stylePromptPhrase: String {
        var details: [String] = []

        if let role {
            details.append(role)
        }

        if let playingStyle {
            details.append(playingStyle)
        }

        let timingPhrase = placements
            .map(\.stylePromptPhrase)
            .filter { !$0.isEmpty }
            .joined(separator: "; ")

        if !timingPhrase.isEmpty {
            details.append(timingPhrase)
        }

        guard !details.isEmpty else {
            return displayName
        }

        return "\(displayName) (\(details.joined(separator: ", ")))"
    }
}

public struct SelectedInstrumentPlacementContext: Codable, Equatable, Sendable {
    public var sectionId: String?
    public var sectionTitle: String?
    public var sectionType: SongSectionType?
    public var startTime: TimeInterval?
    public var endTime: TimeInterval?
    public var direction: String?

    public init(
        sectionId: String? = nil,
        sectionTitle: String? = nil,
        sectionType: SongSectionType? = nil,
        startTime: TimeInterval? = nil,
        endTime: TimeInterval? = nil,
        direction: String? = nil
    ) {
        self.sectionId = sectionId
        self.sectionTitle = sectionTitle
        self.sectionType = sectionType
        self.startTime = startTime
        self.endTime = endTime
        self.direction = direction
    }

    public var stylePromptPhrase: String {
        var parts: [String] = []

        if let sectionTitle, !sectionTitle.isEmpty {
            parts.append(sectionTitle)
        } else if let sectionType {
            parts.append(sectionType.rawValue)
        }

        if let timing = timingPhrase {
            parts.append(timing)
        }

        if let direction, !direction.isEmpty {
            parts.append(direction)
        }

        return parts.joined(separator: " ")
    }

    private var timingPhrase: String? {
        switch (startTime, endTime) {
        case let (start?, end?):
            return "\(Self.formattedTime(start))-\(Self.formattedTime(end))"
        case let (start?, nil):
            return "from \(Self.formattedTime(start))"
        case let (nil, end?):
            return "until \(Self.formattedTime(end))"
        case (nil, nil):
            return nil
        }
    }

    private static func formattedTime(_ time: TimeInterval) -> String {
        let roundedTime = time.rounded()
        if roundedTime == time {
            return "\(Int(roundedTime))s"
        }
        return String(format: "%.1fs", time)
    }
}