package com.suno.android.common_data.use_case

import com.suno.android.common_data.generation.LyricChunk
import javax.inject.Inject
import kotlin.random.Random
import kotlin.time.Duration
import kotlin.time.Duration.Companion.milliseconds

class ProcessLyricsForDisplayUseCase @Inject constructor() {

    private val sectionMarkerRegex = Regex("""\[(\w+(?:\s*\d*)?)\]\s*""")
    private val speakerFindRegex = Regex("""\[([^\]]*?:.*?)\]""")
    private val initialDescRegex = Regex("""^\s*\[.{15,}?\]""")
    private val artistNameRegex = Regex("""\[[^\]]*?:\s*(.*?)\s*\]""")

    private val anonymizedSpeakerNames = (0..99).map { it.toString() }

    private fun makeUniqueList(
        list: List<String>,
    ): List<String> {
        val seen = mutableSetOf<String>()
        return list.filter { seen.add(it) }
    }

    private fun getSpeakerReplMap(
        text: String,
    ): Map<String, String> {
        val matches = artistNameRegex.findAll(text)

        var artistNames = matches.mapNotNull { it.groupValues[1] }.toList()
        artistNames = makeUniqueList(artistNames)
        artistNames = artistNames.sortedBy { it.length }

        val map = mutableMapOf<String, String>()
        for (name in artistNames) {
            if (map.containsKey(name)) continue

            var newName: String? = null
            for ((k, v) in map) {
                if (name.contains(k) && name.replace(k, "").isNotEmpty()) {
                    val speakers = listOf(anonymizedSpeakerNames.random(), v).shuffled()
                    val joiner = if (Random.nextBoolean()) " & " else " and "
                    newName = speakers.joinToString(joiner)
                    break
                }
            }

            map[name] = newName ?: anonymizedSpeakerNames.random()
        }
        return map
    }

    fun anonymizeSpeakers(
        input: String,
    ): String {
        var text = input
        val replMap = getSpeakerReplMap(text)

        val matches = speakerFindRegex.findAll(text).toList().reversed()

        for (match in matches) {
            val fullMatch = match.value
            val innerMatch = match.groupValues[1]

            val parts = innerMatch.split(":", limit = 2)
            if (parts.size == 2) {
                val role = parts[0].trim()
                val name = parts[1].trim()
                val anon = replMap[name] ?: ""
                text = text.replace(fullMatch, "[$role: $anon]")
            }
        }

        text = initialDescRegex.replace(text, "")
        return text.trim()
    }

    fun removeSpeakers(
        input: String,
    ): String {
        var text = input

        text = Regex("""\[([^\]]*?):.*?\]""").replace(text, "[$1]")
        text = initialDescRegex.replace(text, "")

        return text.trim()
    }

    data class WordToken(
        val text: String,
        var startTime: Duration,
        var endTime: Duration,
    )

    data class LyricsLine(
        val text: String,
        val startTime: Duration,
        val endTime: Duration,
        val section: String,
        val words: List<WordToken>,
    )

    private sealed class TokenType {
        object Word : TokenType()
        object Newline : TokenType()
    }

    private data class Token(
        val type: TokenType,
        val text: String,
        val start: Duration,
        val end: Duration,
        val section: String,
    )

    operator fun invoke(
        aligned: List<LyricChunk>,
    ): List<LyricsLine> {
        val tokens = mutableListOf<Token>()
        var currentSection = ""

        for (entry in aligned) {
            var raw = entry.word
            val start = entry.startTime
            var end = if (entry.endTime <= start) start + 50.milliseconds else entry.endTime

            sectionMarkerRegex.find(raw)?.let { match ->
                currentSection = match.groupValues[1]
                raw = sectionMarkerRegex.replace(raw, "")
            }

            if (raw.isEmpty()) continue

            val parts = raw.split("\n")

            for ((i, part) in parts.withIndex()) {
                if (i > 0) {
                    tokens.add(
                        Token(
                            type = TokenType.Newline,
                            text = "",
                            start = start,
                            end = start + 1.milliseconds,
                            section = currentSection,
                        ),
                    )
                }
                if (part.isEmpty()) continue

                tokens.add(
                    Token(
                        type = TokenType.Word,
                        text = part,
                        start = start,
                        end = end,
                        section = currentSection,
                    ),
                )
            }
        }

        data class LineBuild(
            val tokens: List<Token>,
            val section: String,
        )

        val lines = mutableListOf<LineBuild>()
        val buffer = mutableListOf<Token>()

        fun flush() {
            if (buffer.isNotEmpty()) {
                lines.add(LineBuild(buffer.toList(), buffer[0].section))
                buffer.clear()
            }
        }

        for (tok in tokens) {
            when (tok.type) {
                TokenType.Word -> buffer.add(tok)
                TokenType.Newline -> flush()
            }
        }
        flush()

        return lines.mapNotNull { line ->
            val lineText = line.tokens.joinToString("") { it.text }
            val wordsInDisplay = lineText.split(Regex("\\s+"))

            if (wordsInDisplay.isEmpty()) return@mapNotNull null

            val wordTokens = mutableListOf<WordToken>()
            var wordSearchCursor = 0

            for (displayWord in wordsInDisplay) {
                val wordIndex = lineText.indexOf(displayWord, wordSearchCursor)
                if (wordIndex == -1) continue

                wordSearchCursor = wordIndex + displayWord.length

                var wordStart: Duration? = null
                var wordEnd: Duration? = null

                var charPos = 0
                for (tok in line.tokens) {
                    val tokStart = charPos
                    val tokEnd = charPos + tok.text.length

                    val overlaps = !(wordIndex + displayWord.length <= tokStart || wordIndex >= tokEnd)

                    if (overlaps) {
                        wordStart = minOf(wordStart ?: tok.start, tok.start)
                        wordEnd = maxOf(wordEnd ?: tok.end, tok.end)
                    }
                    charPos = tokEnd
                }

                val startTime = wordStart ?: line.tokens.first().start
                val endTime = wordEnd ?: line.tokens.last().end

                wordTokens.add(
                    WordToken(
                        text = displayWord,
                        startTime = startTime,
                        endTime = endTime,
                    ),
                )
            }

            for (i in 1 until wordTokens.size) {
                if (wordTokens[i].startTime < wordTokens[i - 1].endTime) {
                    wordTokens[i].startTime = wordTokens[i - 1].endTime
                }
                if (wordTokens[i].endTime <= wordTokens[i].startTime) {
                    wordTokens[i].endTime = wordTokens[i].startTime + 50.milliseconds
                }
            }

            LyricsLine(
                text = lineText.trim(),
                startTime = wordTokens.first().startTime,
                endTime = wordTokens.last().endTime,
                section = line.section,
                words = wordTokens,
            )
        }
    }
}