package com.suno.android.common_data.generation

import com.suno.android.common_core_utils.Id
import com.suno.android.common_data.mappers.clips.ClipStatus
import com.suno.android.common_data.mappers.clips.SongListData
import kotlinx.collections.immutable.ImmutableMap
import kotlinx.collections.immutable.persistentMapOf
import kotlinx.collections.immutable.toImmutableMap
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asSharedFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.flow.update
import javax.inject.Inject

interface SongGenerationStateStore {
    fun songGenerationStateFlow(): StateFlow<SongGenerationStoreImpl.SongGenerationState>

    fun songGenerationEffectsFlow(): Flow<SongGenerationStoreImpl.SongGenerationEffects>

    fun upsertClips(
        clipStatuses: Map<Id<SongListData>, ClipStatus>,
    )

    fun removeClips(
        clipIds: Set<Id<SongListData>>,
    )
}

class SongGenerationStoreImpl @Inject constructor() : SongGenerationStateStore {
    private val _songGenerationStateFlow = MutableStateFlow(SongGenerationState())

    override fun songGenerationStateFlow(): StateFlow<SongGenerationState> = _songGenerationStateFlow.asStateFlow()

    private val _songGenerationEffects = MutableSharedFlow<SongGenerationEffects>()

    override fun songGenerationEffectsFlow(): Flow<SongGenerationEffects> = _songGenerationEffects.asSharedFlow()

    override fun upsertClips(
        clipStatuses: Map<Id<SongListData>, ClipStatus>,
    ) {
        if (clipStatuses.isEmpty()) return

        _songGenerationStateFlow.update { oldState ->
            oldState.copy(
                clipStatuses = (oldState.clipStatuses + clipStatuses).toImmutableMap(),
            )
        }
    }

    override fun removeClips(
        clipIds: Set<Id<SongListData>>,
    ) {
        if (clipIds.isEmpty()) return

        _songGenerationStateFlow.update { oldState ->
            oldState.copy(
                clipStatuses = (oldState.clipStatuses - clipIds).toImmutableMap(),
            )
        }
    }

    data class SongGenerationState(
        val clipStatuses: ImmutableMap<Id<SongListData>, ClipStatus> = persistentMapOf(),
    ) {
        val generatingClipIds: Set<Id<SongListData>> by lazy { clipStatuses.filterValues { it.isGenerating }.keys }
        val readyClipIds: Set<Id<SongListData>> by lazy { clipStatuses.filterValues { it.isReady }.keys }
        val nonTerminalClipIds: Set<Id<SongListData>> by lazy { clipStatuses.filterValues { !it.isTerminal }.keys }
    }

    sealed interface SongGenerationEffects {
        data class Error(
            val reason: String,
        ) : SongGenerationEffects
    }
}