package com.suno.android.common_analytics.orpheus

import com.segment.analytics.kotlin.core.Analytics
import com.suno.android.common_analytics.APP_EVENT_NAME
import com.suno.android.common_analytics.ElementType
import com.suno.android.common_analytics.PLATFORM
import com.suno.android.common_core_utils.ApplicationCoroutineScope
import com.suno.android.common_core_utils.BuildConfig
import com.suno.android.common_core_utils.DispatcherIO
import com.suno.android.common_core_utils.Id
import com.suno.android.common_core_utils.SunoLogger
import com.suno.android.common_data.repos.orpheus.OrpheusSessionStore
import com.suno.android.common_data.repos.orpheus.models.OrpheusMessage
import com.suno.android.common_data.repos.orpheus.models.OrpheusSession
import com.suno.android.common_data.user.UserSessionRepository
import kotlinx.coroutines.CoroutineDispatcher
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.catch
import kotlinx.coroutines.flow.consumeAsFlow
import kotlinx.coroutines.flow.distinctUntilChanged
import kotlinx.coroutines.flow.filterNotNull
import kotlinx.coroutines.flow.launchIn
import kotlinx.coroutines.flow.map
import kotlinx.coroutines.flow.onEach
import kotlinx.coroutines.launch
import kotlinx.serialization.encodeToString
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.encodeToJsonElement
import java.time.Instant
import java.util.concurrent.atomic.AtomicReference
import javax.inject.Inject
import javax.inject.Singleton

/**
 * Manages analytics tracking for Orpheus events.
 */
interface OrpheusAnalyticsManager {
    fun trackSessionOpened()

    fun trackSessionClosed()

    fun trackSwitchToCustomMode()

    fun trackSwitchToOrpheusMode()

    /**
     * Tracks when a user sends a message.
     *
     * @param messageId The user's message ID
     */
    fun trackMessageSent(
        messageId: Id<OrpheusMessage>,
    )

    /**
     * Tracks when the first chunk of an Orpheus response is received.
     *
     * @param messageId The assistant's message ID
     */
    fun trackResponseBegin(
        messageId: Id<OrpheusMessage>,
    )

    /**
     * Tracks when an Orpheus response is fully received.
     *
     * @param messageId The assistant's message ID
     */
    fun trackResponseEnd(
        messageId: Id<OrpheusMessage>,
    )

    /**
     * Gets the current Orpheus session ID for cross-referencing with other analytics.
     */
    fun getCurrentSessionId(): Id<OrpheusAnalyticsSession>?

    fun analyticDebugLogTailFlow(): Flow<String>
}

private const val ANALYTICS_CHANNEL_CAPACITY = 1000

@Singleton
class OrpheusAnalyticsManagerImpl @Inject constructor(
    loggerFactory: SunoLogger.Factory,
    @ApplicationCoroutineScope private val applicationScope: CoroutineScope,
    @DispatcherIO private val dispatcherIO: CoroutineDispatcher,
    private val analytics: Analytics,
    private val userSessionRepository: UserSessionRepository,
    private val sessionStore: OrpheusSessionStore,
) : OrpheusAnalyticsManager {
    private val logger = loggerFactory.create(this@OrpheusAnalyticsManagerImpl)
    private val currentAnalyticsSession = AtomicReference<OrpheusAnalyticsSession>(null)
    private val currentUserId = AtomicReference<String>(null)
    private val analyticsChannel = Channel<QueuedOrpheusAnalyticsEvent>(ANALYTICS_CHANNEL_CAPACITY)
    private val _analyticDebugLogTailFlow = MutableSharedFlow<String>()
    private val prettyJson by lazy { Json { prettyPrint = true } }

    private val currentChatId: Id<OrpheusSession>? get() = sessionStore.sessionFlow.value.dataOrNull?.sessionId
    private val lastUserMessageId = AtomicReference<Id<OrpheusMessage>?>(null)

    init {
        startBackgroundProcessor()
        observeCurrentUserId()
    }

    private fun startBackgroundProcessor() {
        applicationScope.launch(dispatcherIO) {
            analyticsChannel.consumeAsFlow().collect { queuedEvent ->
                runCatching {
                    processAnalyticsEvent(queuedEvent)
                }.onFailure {
                    logger.e(it) { "Failed to process analytics event" }
                }
            }
        }
    }

    private suspend fun processAnalyticsEvent(
        queuedEvent: QueuedOrpheusAnalyticsEvent,
    ) {
        val userId = currentUserId.get() ?: return

        val event = OrpheusAnalyticsEvent(
            actionName = queuedEvent.action.actionName,
            userId = userId,
            elementType = ElementType.Orpheus,
            metadata = queuedEvent.metadata,
        )
        val jsonElement = Json.encodeToJsonElement(event) as? JsonObject ?: return

        analytics.track(
            name = APP_EVENT_NAME,
            properties = jsonElement,
        )

        if (BuildConfig.IS_STAFF) {
            val prettyJsonString = prettyJson.encodeToString(jsonElement)
            _analyticDebugLogTailFlow.emit(prettyJsonString)
        }

        logger.d {
            "processAnalyticsEvent (thread: ${Thread.currentThread().name}) - ${queuedEvent.action.actionName}\n" +
                prettyJson.encodeToString(jsonElement)
        }
    }

    private fun observeCurrentUserId() {
        userSessionRepository.sessionConfigurationStateFlow()
            .map { it.user }
            .filterNotNull()
            .distinctUntilChanged()
            .onEach { user ->
                user.id?.let { userId ->
                    analytics.identify(userId = userId)
                    currentUserId.set(userId)
                }
            }
            .catch { logger.e(it) }
            .launchIn(applicationScope)
    }

    override fun trackSessionOpened() {
        startNewSession()
        queueAnalyticsEvent(
            action = OrpheusAnalyticsAction.CreateSessionOpened,
            metadata = createMetadata(chatId = currentChatId),
        )
    }

    override fun trackSessionClosed() {
        queueAnalyticsEvent(
            action = OrpheusAnalyticsAction.CreateSessionClosed,
            metadata = createMetadata(chatId = currentChatId),
        )
        clearSession()
    }

    override fun trackSwitchToCustomMode() {
        queueAnalyticsEvent(
            action = OrpheusAnalyticsAction.SwitchFromOrpheusToCustomMode,
            metadata = createMetadata(chatId = currentChatId),
        )
    }

    override fun trackSwitchToOrpheusMode() {
        queueAnalyticsEvent(
            action = OrpheusAnalyticsAction.SwitchFromCustomToOrpheusMode,
            metadata = createMetadata(chatId = currentChatId),
        )
    }

    override fun trackMessageSent(
        messageId: Id<OrpheusMessage>,
    ) {
        queueAnalyticsEvent(
            action = OrpheusAnalyticsAction.OrpheusMessageSendTime,
            metadata = createMetadata(
                chatId = currentChatId,
                messageId = messageId,
            ),
        )
        lastUserMessageId.set(messageId)
    }

    override fun trackResponseBegin(
        messageId: Id<OrpheusMessage>,
    ) {
        queueAnalyticsEvent(
            action = OrpheusAnalyticsAction.OrpheusMessageBeginResponseTime,
            metadata = createMetadata(
                chatId = currentChatId,
                messageId = messageId,
                precedingMessageId = lastUserMessageId.get(),
            ),
        )
    }

    override fun trackResponseEnd(
        messageId: Id<OrpheusMessage>,
    ) {
        queueAnalyticsEvent(
            action = OrpheusAnalyticsAction.OrpheusMessageEndResponseTime,
            metadata = createMetadata(
                chatId = currentChatId,
                messageId = messageId,
                precedingMessageId = lastUserMessageId.get(),
            ),
        )
    }

    override fun getCurrentSessionId(): Id<OrpheusAnalyticsSession>? = currentAnalyticsSession.get()?.sessionId

    override fun analyticDebugLogTailFlow(): Flow<String> = _analyticDebugLogTailFlow

    private fun startNewSession() {
        val newSession = OrpheusAnalyticsSession(sessionId = Id())
        currentAnalyticsSession.set(newSession)
        logger.d { "Auto-generated Orpheus session: $newSession" }
    }

    private fun clearSession() {
        currentAnalyticsSession.set(null)
    }

    private fun createMetadata(
        chatId: Id<OrpheusSession>? = null,
        messageId: Id<OrpheusMessage>? = null,
        precedingMessageId: Id<OrpheusMessage>? = null,
    ): OrpheusAnalyticsMetadata {
        val sessionId = currentAnalyticsSession.get()?.sessionId
        return buildMap {
            put("session_id", sessionId?.value)
            put("timestamp", Instant.now().toString())
            put("platform", PLATFORM)

            chatId?.let { put("chat_id", it.value) }
            messageId?.let { put("message_id", it.value) }
            precedingMessageId?.let { put("preceding_message_id", it.value) }
        }
    }

    private fun queueAnalyticsEvent(
        action: OrpheusAnalyticsAction,
        metadata: OrpheusAnalyticsMetadata,
    ) {
        val queuedEvent = QueuedOrpheusAnalyticsEvent(
            action = action,
            metadata = metadata,
        )
        val result = analyticsChannel.trySend(queuedEvent)
        if (!result.isSuccess) {
            logger.w { "Failed to queue analytics event: ${result.exceptionOrNull()}" }
        }
    }
}

private data class QueuedOrpheusAnalyticsEvent(
    val action: OrpheusAnalyticsAction,
    val metadata: OrpheusAnalyticsMetadata,
)