package com.suno.android.common_analytics

import com.segment.analytics.kotlin.core.Analytics
import com.suno.android.common_analytics.orpheus.OrpheusAnalyticsManagerImpl
import com.suno.android.common_core_utils.Id
import com.suno.android.common_core_utils.SunoLogger
import com.suno.android.common_core_utils.model.AsyncData
import com.suno.android.common_data.repos.orpheus.OrpheusSessionError
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.UserSessionConfiguration
import com.suno.android.common_data.user.UserSessionRepository
import com.suno.android.common_networking.remote.session.User
import io.mockk.clearMocks
import io.mockk.every
import io.mockk.mockk
import io.mockk.slot
import io.mockk.verify
import kotlinx.collections.immutable.persistentListOf
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.test.StandardTestDispatcher
import kotlinx.coroutines.test.TestScope
import kotlinx.coroutines.test.advanceUntilIdle
import kotlinx.coroutines.test.runTest
import kotlinx.serialization.json.JsonObject
import org.junit.Assert.assertEquals
import org.junit.Assert.assertNotNull
import org.junit.Assert.assertNull
import org.junit.Assert.assertTrue
import org.junit.Before
import org.junit.Test
import kotlin.time.Duration.Companion.seconds

private const val TEST_USER_ID = "test-user-123"
private const val TEST_CHAT_ID = "chat-939"

@Suppress("ktlint:standard:max-line-length")
@OptIn(ExperimentalCoroutinesApi::class)
class OrpheusAnalyticsManagerImplTest {
    private val analytics = mockk<Analytics>(relaxed = true)
    private val userSessionRepository = mockk<UserSessionRepository>()
    private val sessionStore = mockk<OrpheusSessionStore>()
    private val loggerFactory = mockk<SunoLogger.Factory>(relaxed = true)

    private val testDispatcher = StandardTestDispatcher()
    private val applicationScope = TestScope(testDispatcher)

    private lateinit var sessionConfigFlow: MutableStateFlow<UserSessionConfiguration>
    private lateinit var sessionStoreFlow: MutableStateFlow<AsyncData<OrpheusSessionError, OrpheusSession?>>

    private lateinit var subject: OrpheusAnalyticsManagerImpl

    @Before
    fun setup() {
        sessionConfigFlow = MutableStateFlow(createUserSessionConfiguration())
        sessionStoreFlow = MutableStateFlow(AsyncData.Ready(createOrpheusSession()))

        every { userSessionRepository.sessionConfigurationStateFlow() } returns sessionConfigFlow.asStateFlow()
        every { sessionStore.sessionFlow } returns sessionStoreFlow.asStateFlow()

        subject = OrpheusAnalyticsManagerImpl(
            loggerFactory = loggerFactory,
            applicationScope = applicationScope,
            dispatcherIO = testDispatcher,
            analytics = analytics,
            userSessionRepository = userSessionRepository,
            sessionStore = sessionStore,
        )

        testDispatcher.scheduler.advanceUntilIdle()
        clearMocks(analytics, answers = false)
    }

    @Test
    fun `given no active orpheus session when trackSessionOpened called then tracks createSessionOpened event`() =
        runTest(testDispatcher) {
            val propertiesSlot = slot<JsonObject>()

            subject.trackSessionOpened()
            advanceUntilIdle()

            verify(exactly = 1) {
                analytics.track(
                    name = APP_EVENT_NAME,
                    properties = capture(propertiesSlot),
                )
            }
            propertiesSlot.captured.assertOrpheusEvent(
                expectedActionName = "createSessionOpened",
                expectedChatId = TEST_CHAT_ID,
            )
            assertNotNull(subject.getCurrentSessionId())
        }

    @Test
    fun `given orpheus session opened when trackSessionClosed called then tracks createSessionClosed event`() =
        runTest(testDispatcher) {
            subject.trackSessionOpened()
            advanceUntilIdle()
            clearMocks(analytics, answers = false)
            val propertiesSlot = slot<JsonObject>()

            subject.trackSessionClosed()
            advanceUntilIdle()

            verify(exactly = 1) {
                analytics.track(
                    name = APP_EVENT_NAME,
                    properties = capture(propertiesSlot),
                )
            }
            propertiesSlot.captured.assertOrpheusEvent(
                expectedActionName = "createSessionClosed",
                expectedChatId = TEST_CHAT_ID,
            )
            assertNull(subject.getCurrentSessionId())
        }

    @Test
    fun `given orpheus session opened when trackSwitchToCustomMode called then tracks mode switch event`() =
        runTest(testDispatcher) {
            subject.trackSessionOpened()
            advanceUntilIdle()
            clearMocks(analytics, answers = false)
            val propertiesSlot = slot<JsonObject>()

            subject.trackSwitchToCustomMode()
            advanceUntilIdle()

            verify(exactly = 1) {
                analytics.track(
                    name = APP_EVENT_NAME,
                    properties = capture(propertiesSlot),
                )
            }
            propertiesSlot.captured.assertOrpheusEvent(
                expectedActionName = "switchFromOrpheusToCustomMode",
                expectedChatId = TEST_CHAT_ID,
            )
        }

    @Test
    fun `given orpheus session opened when trackSwitchToOrpheusMode called then tracks mode switch event`() =
        runTest(testDispatcher) {
            subject.trackSessionOpened()
            advanceUntilIdle()
            clearMocks(analytics, answers = false)
            val propertiesSlot = slot<JsonObject>()

            subject.trackSwitchToOrpheusMode()
            advanceUntilIdle()

            verify(exactly = 1) {
                analytics.track(
                    name = APP_EVENT_NAME,
                    properties = capture(propertiesSlot),
                )
            }
            propertiesSlot.captured.assertOrpheusEvent(
                expectedActionName = "switchFromCustomToOrpheusMode",
                expectedChatId = TEST_CHAT_ID,
            )
        }

    @Test
    fun `given orpheus session with no prior messages when trackMessageSent called then tracks message send event`() =
        runTest(testDispatcher) {
            subject.trackSessionOpened()
            advanceUntilIdle()
            clearMocks(analytics, answers = false)
            val messageId = Id<OrpheusMessage>("0000")
            val propertiesSlot = slot<JsonObject>()

            subject.trackMessageSent(messageId)
            advanceUntilIdle()

            verify(exactly = 1) {
                analytics.track(
                    name = APP_EVENT_NAME,
                    properties = capture(propertiesSlot),
                )
            }
            propertiesSlot.captured.assertOrpheusEvent(
                expectedActionName = "orpheusMessageSendTime",
                expectedChatId = TEST_CHAT_ID,
                expectedMessageId = messageId.value,
            )
        }

    @Test
    fun `given orpheus session when streaming assistant response tracked called then tracks response events`() =
        runTest(testDispatcher) {
            subject.trackSessionOpened()
            advanceUntilIdle()
            clearMocks(analytics, answers = false)
            val userMessageId = Id<OrpheusMessage>("0000")
            val assistantMessageId = Id<OrpheusMessage>("1111")
            val propertiesSlot = mutableListOf<JsonObject>()

            subject.trackMessageSent(userMessageId)
            subject.trackResponseBegin(assistantMessageId)
            subject.trackResponseEnd(assistantMessageId)
            advanceUntilIdle()

            verify(exactly = 3) {
                analytics.track(
                    name = APP_EVENT_NAME,
                    properties = capture(propertiesSlot),
                )
            }
            propertiesSlot[0].assertOrpheusEvent(
                expectedActionName = "orpheusMessageSendTime",
                expectedChatId = TEST_CHAT_ID,
                expectedMessageId = userMessageId.value,
            )
            propertiesSlot[1].assertOrpheusEvent(
                expectedActionName = "orpheusMessageBeginResponseTime",
                expectedChatId = TEST_CHAT_ID,
                expectedMessageId = assistantMessageId.value,
                expectedPrecedingMessageId = userMessageId.value,
            )
            propertiesSlot[2].assertOrpheusEvent(
                expectedActionName = "orpheusMessageEndResponseTime",
                expectedChatId = TEST_CHAT_ID,
                expectedMessageId = assistantMessageId.value,
                expectedPrecedingMessageId = userMessageId.value,
            )
        }

    @Test
    fun `given orpheus session when multiple streaming assistant response tracked called then tracks multiple response events in order`() =
        runTest(testDispatcher) {
            subject.trackSessionOpened()
            advanceUntilIdle()
            clearMocks(analytics, answers = false)
            val userMessageId = Id<OrpheusMessage>("0000")
            val assistantMessageId = Id<OrpheusMessage>("1111")
            val assistantSongGenMessageId = Id<OrpheusMessage>("2222")
            val propertiesSlot = mutableListOf<JsonObject>()

            subject.trackMessageSent(userMessageId)
            subject.trackResponseBegin(assistantSongGenMessageId)
            subject.trackResponseBegin(assistantMessageId)
            subject.trackResponseEnd(assistantMessageId)
            subject.trackResponseEnd(assistantSongGenMessageId)
            advanceUntilIdle()

            verify(exactly = 5) {
                analytics.track(
                    name = APP_EVENT_NAME,
                    properties = capture(propertiesSlot),
                )
            }
            propertiesSlot[0].assertOrpheusEvent(
                expectedActionName = "orpheusMessageSendTime",
                expectedChatId = TEST_CHAT_ID,
                expectedMessageId = userMessageId.value,
            )
            propertiesSlot[1].assertOrpheusEvent(
                expectedActionName = "orpheusMessageBeginResponseTime",
                expectedChatId = TEST_CHAT_ID,
                expectedMessageId = assistantSongGenMessageId.value,
                expectedPrecedingMessageId = userMessageId.value,
            )
            propertiesSlot[2].assertOrpheusEvent(
                expectedActionName = "orpheusMessageBeginResponseTime",
                expectedChatId = TEST_CHAT_ID,
                expectedMessageId = assistantMessageId.value,
                expectedPrecedingMessageId = userMessageId.value,
            )
            propertiesSlot[3].assertOrpheusEvent(
                expectedActionName = "orpheusMessageEndResponseTime",
                expectedChatId = TEST_CHAT_ID,
                expectedMessageId = assistantMessageId.value,
                expectedPrecedingMessageId = userMessageId.value,
            )
            propertiesSlot[4].assertOrpheusEvent(
                expectedActionName = "orpheusMessageEndResponseTime",
                expectedChatId = TEST_CHAT_ID,
                expectedMessageId = assistantSongGenMessageId.value,
                expectedPrecedingMessageId = userMessageId.value,
            )
        }

    @Test
    fun `given orpheus session opened when getCurrentSessionId called then returns session ID`() =
        runTest(testDispatcher) {
            subject.trackSessionOpened()
            advanceUntilIdle()

            val sessionId = subject.getCurrentSessionId()

            assertNotNull(sessionId)
        }

    @Test
    fun `given no active session when tracking event then still tracks event with null sessionId`() =
        runTest(testDispatcher) {
            // Don't open session!

            val messageId = Id<OrpheusMessage>("0000")
            val propertiesSlot = mutableListOf<JsonObject>()

            subject.trackMessageSent(messageId)
            advanceUntilIdle()

            verify(exactly = 1) {
                analytics.track(
                    name = APP_EVENT_NAME,
                    properties = capture(propertiesSlot),
                )
            }
            propertiesSlot[0].assertOrpheusEvent(
                expectedActionName = "orpheusMessageSendTime",
                expectedChatId = TEST_CHAT_ID,
                expectedMessageId = messageId.value,
            )
            val context = propertiesSlot[0]["context"] as JsonObject
            assertTrue(context["session_id"]?.toString() == "null")
        }

    @Test
    fun `given null precedingMessageId when tracking response then event sent without it`() =
        runTest(testDispatcher) {
            subject.trackSessionOpened()
            advanceUntilIdle()
            clearMocks(analytics, answers = false)

            val messageId = Id<OrpheusMessage>("0000")
            val propertiesSlot = slot<JsonObject>()

            subject.trackResponseBegin(messageId)
            advanceUntilIdle()

            verify(exactly = 1) {
                analytics.track(
                    name = APP_EVENT_NAME,
                    properties = capture(propertiesSlot),
                )
            }
            propertiesSlot.captured.assertOrpheusEvent(
                expectedActionName = "orpheusMessageBeginResponseTime",
                expectedChatId = TEST_CHAT_ID,
                expectedMessageId = messageId.value,
                expectedPrecedingMessageId = null,
            )
        }

    @Test
    fun `given multiple events tracked when processed then all events sent`() =
        runTest(testDispatcher) {
            subject.trackSessionOpened()
            advanceUntilIdle()
            clearMocks(analytics, answers = false)
            val propertiesSlots = mutableListOf<JsonObject>()

            subject.trackMessageSent(Id("0000"))
            subject.trackResponseBegin(Id("1111"))
            subject.trackResponseEnd(Id("2222"))
            advanceUntilIdle()

            verify(exactly = 3) {
                analytics.track(
                    name = APP_EVENT_NAME,
                    properties = capture(propertiesSlots),
                )
            }
            assertEquals(3, propertiesSlots.size)
        }

    @Test
    fun `given session closed when getCurrentSessionId called then returns null`() =
        runTest(testDispatcher) {
            subject.trackSessionOpened()
            advanceUntilIdle()
            subject.trackSessionClosed()
            advanceUntilIdle()

            val sessionId = subject.getCurrentSessionId()

            assertNull(sessionId)
        }

    // Helper Methods

    private fun createUserSessionConfiguration(
        userId: String? = TEST_USER_ID,
    ): UserSessionConfiguration =
        UserSessionConfiguration(
            user = User(
                id = userId,
                handle = "@testuser",
                avatarImageUrl = "https://example.com/avatar.jpg",
                displayName = "Test User",
                email = "test@example.com",
                isHandleUpdated = true,
                profileDescription = "Test bio",
                username = "testuser",
            ),
            roles = null,
            statsigCustomProperties = null,
        )

    private fun createOrpheusSession(
        sessionId: String = TEST_CHAT_ID,
    ): OrpheusSession =
        OrpheusSession(
            sessionId = Id(sessionId),
            createdAt = 0.seconds,
            currentModel = "test-model",
            messages = persistentListOf(),
        )

    private fun JsonObject.assertOrpheusEvent(
        expectedActionName: String,
        expectedChatId: String? = null,
        expectedMessageId: String? = null,
        expectedPrecedingMessageId: String? = null,
    ) {
        assertEquals(expectedActionName, this["actionName"].toString().removeSurrounding("\""))
        assertEquals(TEST_USER_ID, this["userId"].toString().removeSurrounding("\""))
        assertEquals("orpheus", this["elementType"].toString().removeSurrounding("\""))

        val context = this["context"] as JsonObject
        assertEquals("android", context["platform"].toString().removeSurrounding("\""))
        assertTrue(context.containsKey("session_id"))
        assertTrue(context.containsKey("timestamp"))

        if (expectedChatId != null) {
            assertEquals(expectedChatId, context["chat_id"].toString().removeSurrounding("\""))
        }

        if (expectedMessageId != null) {
            assertEquals(expectedMessageId, context["message_id"].toString().removeSurrounding("\""))
        }

        if (expectedPrecedingMessageId != null) {
            assertEquals(
                expectedPrecedingMessageId,
                context["preceding_message_id"].toString().removeSurrounding("\""),
            )
        } else if (this@assertOrpheusEvent["actionName"].toString().contains("Response")) {
            // Response events with null preceding_message_id should not include the key at all
            assertTrue(!context.containsKey("preceding_message_id"))
        }
    }
}