music-assistant/Tests/MusicAssistantCoreTests/SongProjectPromptCompilerTests.swift

import MusicAssistantCore
import XCTest

final class SongProjectPromptCompilerTests: XCTestCase {
    func testCompilerUsesSelectedInstrumentsInSunoStylePrompt() throws {
        let project = SongProject(
            title: "Compiled Song",
            idea: "Prepare Suno output",
            genres: [
                GenreStyle(id: "arabic-pop", name: "Arabic Pop"),
                GenreStyle(id: "cinematic", name: "Cinematic")
            ],
            moods: [MoodTag(id: "hopeful", name: "Hopeful")],
            bpm: ManualAutoValue(mode: .manual, value: 96),
            maqam: ManualAutoValue(mode: .manual, value: "Hijaz"),
            sections: [
                SongSection(id: "chorus", type: .chorus, title: "Chorus", lyrics: "Section chorus")
            ],
            instruments: [
                InstrumentTrack(
                    instrumentId: "oud",
                    selected: true,
                    variant: "Arabic Oud",
                    playingStyle: "picked",
                    role: "main hook",
                    placements: [InstrumentPlacement(sectionId: "chorus", direction: "answer the vocal")]
                ),
                InstrumentTrack(
                    instrumentId: "drum-kit",
                    selected: false,
                    variant: "Drum Kit"
                )
            ],
            vocalists: [Vocalist(id: "lead", label: "Lead", voiceType: "warm tenor")],
            lyrics: Lyrics(text: "Approved lyrics"),
            productionDirections: [ProductionDirection(id: "lift", text: "wide chorus lift")]
        )

        let output = try SongProjectPromptCompiler().compile(project: project)

        XCTAssertEqual(output.lyricsText, "Approved lyrics")
        XCTAssertTrue(output.stylePrompt.contains("Arabic Pop + Cinematic"))
        XCTAssertTrue(output.stylePrompt.contains("Arabic Oud (main hook, picked, Chorus answer the vocal)"))
        XCTAssertTrue(output.stylePrompt.contains("Lead warm tenor"))
        XCTAssertTrue(output.stylePrompt.contains("96 BPM"))
        XCTAssertTrue(output.stylePrompt.contains("maqam Hijaz"))
        XCTAssertTrue(output.stylePrompt.contains("wide chorus lift"))
        XCTAssertFalse(output.stylePrompt.contains("Drum Kit"))
    }

    func testCompilerFallsBackToOrderedSectionLyrics() throws {
        let project = SongProject(
            title: "Section Lyrics",
            idea: "Compile sections",
            sections: [
                SongSection(id: "chorus", type: .chorus, title: "Chorus", order: 1, lyrics: "Hook line"),
                SongSection(id: "verse", type: .verse, title: "Verse", order: 0, lyrics: "Verse line")
            ]
        )

        let output = try SongProjectPromptCompiler().compile(project: project)

        XCTAssertEqual(output.lyricsText, "[Verse]\nVerse line\n\n[Chorus]\nHook line")
    }

    func testCompilerIncludesSectionPerformanceDirectivesInFallbackLyrics() throws {
        let project = SongProject(
            title: "Directed Sections",
            idea: "Compile performance directions",
            sections: [
                SongSection(
                    id: "intro",
                    type: .intro,
                    title: "Intro",
                    lyrics: "First line",
                    vocalDirection: "whispered lead"
                ),
                SongSection(
                    id: "chorus",
                    type: .chorus,
                    title: "Chorus",
                    order: 1,
                    lyrics: "Original hook"
                )
            ],
            lyrics: Lyrics(sectionTexts: ["chorus": "Approved hook"])
        )

        let output = try SongProjectPromptCompiler().compile(project: project)

        XCTAssertEqual(
            output.lyricsText,
            "[Intro - whispered lead]\nFirst line\n\n[Chorus]\nApproved hook"
        )
    }

    func testCompilerIncludesStyleInputsFromTheSongProject() throws {
        let project = SongProject(
            title: "Full Style",
            idea: "Compile every approved style input",
            genres: [
                GenreStyle(id: "folk", name: "Folk"),
                GenreStyle(id: "electronic", name: "Electronic")
            ],
            moods: [MoodTag(id: "warm", name: "Warm")],
            emotionalArc: [
                EmotionalArcPoint(
                    id: "chorus-arc",
                    sectionID: "chorus",
                    emotion: "joyful",
                    energy: 8,
                    direction: "lift into release"
                )
            ],
            bpm: ManualAutoValue(mode: .manual, value: 108),
            key: ManualAutoValue(mode: .manual, value: "D minor"),
            scale: ManualAutoValue(mode: .manual, value: "Dorian"),
            sections: [
                SongSection(
                    id: "verse",
                    type: .verse,
                    title: "Verse",
                    emotion: "yearning",
                    energy: 3,
                    productionDirection: "intimate opening"
                ),
                SongSection(id: "chorus", type: .chorus, title: "Chorus", order: 1)
            ],
            vocalists: [
                Vocalist(id: "lead", label: "Lead", voiceType: "clear alto", performanceStyle: "expressive")
            ],
            productionDirections: [
                ProductionDirection(id: "chorus-width", sectionID: "chorus", text: "wide chorus"),
                ProductionDirection(id: "master", text: "warm tape saturation")
            ]
        )

        let output = try SongProjectPromptCompiler().compile(project: project)

        XCTAssertEqual(
            output.stylePrompt,
            "Folk + Electronic; mood: Warm; Lead clear alto expressive; 108 BPM, key D minor, Dorian scale; emotion: Verse, yearning, energy 3; Chorus, joyful, energy 8, lift into release; production: Verse, intimate opening; Chorus, wide chorus; warm tape saturation"
        )
    }
}