@testable import MusicAssistantCore
import XCTest

final class ApplicationRuleInjectionTests: XCTestCase {
    func testInjectingServiceAddsPrivateRulesToEveryAIRequestType() async throws {
        let baseService = RecordingAIService()
        let service = ApplicationRuleInjectingAIService(
            baseService: baseService,
            ruleProvider: StaticApplicationRuleProvider(
                rules: PrivateApplicationRules(content: "Use private product rules.")
            )
        )
        let project = SongProject(title: "Rule Test", idea: "Test rules")

        _ = try await service.generateSongProject(
            from: SongProjectGenerationRequest(
                context: AIRequestContext(userInstruction: "Generate.")
            )
        )
        _ = try await service.discussSongProject(
            from: SongProjectDiscussionRequest(
                context: AIRequestContext(userInstruction: "Discuss."),
                project: project
            )
        )
        _ = try await service.reviseLyrics(
            from: LyricsRevisionRequest(
                context: AIRequestContext(userInstruction: "Improve."),
                project: project,
                sourceLyrics: Lyrics(text: "Draft"),
                mode: .improve
            )
        )
        _ = try await service.proposeProjectUpdate(
            from: SongProjectUpdateRequest(
                context: AIRequestContext(userInstruction: "Update."),
                project: project,
                allowedScopes: [.lyrics]
            )
        )

        let recordedRuleContents = await baseService.recordedRuleContents()
        XCTAssertEqual(
            recordedRuleContents,
            [
                "Use private product rules.",
                "Use private product rules.",
                "Use private product rules.",
                "Use private product rules."
            ]
        )
    }

    func testInjectingServiceDoesNotAddPrivateRulesToVisibleConversation() async throws {
        let baseService = RecordingAIService()
        let service = ApplicationRuleInjectingAIService(
            baseService: baseService,
            ruleProvider: StaticApplicationRuleProvider(
                rules: PrivateApplicationRules(content: "Hidden owner rules")
            )
        )
        let visibleMessage = AIConversationMessage(role: .user, content: "Visible user text.")

        _ = try await service.generateSongProject(
            from: SongProjectGenerationRequest(
                context: AIRequestContext(
                    userInstruction: "Generate.",
                    conversation: [visibleMessage]
                )
            )
        )

        let recordedConversation = await baseService.recordedConversations().first
        XCTAssertEqual(recordedConversation, [visibleMessage])
    }

    func testBlankPrivateRulesAreIgnored() {
        XCTAssertNil(PrivateApplicationRules(content: "  \n  "))
    }
}

private actor RecordingAIService: AIService {
    private var contexts: [AIRequestContext] = []

    func recordedRuleContents() -> [String?] {
        contexts.map { $0.privateApplicationRules?.content }
    }

    func recordedConversations() -> [[AIConversationMessage]] {
        contexts.map(\.conversation)
    }

    func generateSongProject(from request: SongProjectGenerationRequest) async throws -> SongProjectGenerationResult {
        contexts.append(request.context)
        return SongProjectGenerationResult(
            project: SongProject(title: "Generated", idea: request.context.userInstruction)
        )
    }

    func discussSongProject(from request: SongProjectDiscussionRequest) async throws -> SongProjectDiscussionResult {
        contexts.append(request.context)
        return SongProjectDiscussionResult(questions: ["Which direction should we take?"])
    }

    func reviseLyrics(from request: LyricsRevisionRequest) async throws -> LyricsRevisionResult {
        contexts.append(request.context)
        return LyricsRevisionResult(lyrics: request.sourceLyrics)
    }

    func proposeProjectUpdate(from request: SongProjectUpdateRequest) async throws -> SongProjectUpdateResult {
        contexts.append(request.context)
        return SongProjectUpdateResult(project: request.project)
    }
}