ottobotv2/src/main/kotlin/nl/voidcorp/discord/music/CustomYoutubeSourceManager.kt

package nl.voidcorp.discord.music

import com.sedmelluq.discord.lavaplayer.player.DefaultAudioPlayerManager
import com.sedmelluq.discord.lavaplayer.source.AudioSourceManager
import com.sedmelluq.discord.lavaplayer.source.youtube.YoutubeAudioSourceManager
import com.sedmelluq.discord.lavaplayer.source.youtube.YoutubeAudioTrack
import com.sedmelluq.discord.lavaplayer.source.youtube.YoutubeMixProvider
import com.sedmelluq.discord.lavaplayer.source.youtube.YoutubeSearchProvider
import com.sedmelluq.discord.lavaplayer.tools.DataFormatTools
import com.sedmelluq.discord.lavaplayer.tools.DataFormatTools.convertToMapLayout
import com.sedmelluq.discord.lavaplayer.tools.ExceptionTools
import com.sedmelluq.discord.lavaplayer.tools.FriendlyException
import com.sedmelluq.discord.lavaplayer.tools.JsonBrowser
import com.sedmelluq.discord.lavaplayer.tools.io.HttpClientTools
import com.sedmelluq.discord.lavaplayer.tools.io.HttpConfigurable
import com.sedmelluq.discord.lavaplayer.tools.io.HttpInterface
import com.sedmelluq.discord.lavaplayer.tools.io.HttpInterfaceManager
import com.sedmelluq.discord.lavaplayer.track.*
import org.apache.commons.io.IOUtils
import org.apache.http.client.config.RequestConfig
import org.apache.http.client.methods.HttpGet
import org.apache.http.client.utils.URIBuilder
import org.apache.http.client.utils.URLEncodedUtils
import org.apache.http.impl.client.HttpClientBuilder
import org.jsoup.Jsoup
import org.jsoup.nodes.Document
import org.jsoup.nodes.Element
import org.slf4j.LoggerFactory
import java.io.DataInput
import java.io.DataOutput
import java.io.IOException
import java.net.URISyntaxException
import java.net.URLEncoder
import java.nio.charset.Charset
import java.util.*
import java.util.function.Consumer
import java.util.regex.Pattern

/**
 * Audio source manager that implements finding Youtube videos or playlists based on an URL or ID.
 */
class CustomYoutubeSourceManager : YoutubeAudioSourceManager(), AudioSourceManager, HttpConfigurable {

    private val extractors = arrayOf(
        Extractor(directVideoIdPattern) { id -> loadTrackWithVideoId(id, false) },
        Extractor(
            Pattern.compile("^$PLAYLIST_ID_REGEX$")
        ) { id -> loadPlaylistWithId(id, null) },
        Extractor(
            Pattern.compile("^$PROTOCOL_REGEX$DOMAIN_REGEX/.*")
        ) {
            println(it)
            this.loadFromMainDomain(it) },
        Extractor(
            Pattern.compile("^$PROTOCOL_REGEX$SHORT_DOMAIN_REGEX/.*")
        ) { this.loadFromShortDomain(it) }
    )

    private val httpInterfaceManager: HttpInterfaceManager = HttpClientTools.createDefaultThreadLocalManager()
    private val searchProvider: YoutubeSearchProvider = YoutubeSearchProvider(this)
    private val mixProvider: YoutubeMixProvider = YoutubeMixProvider(this)


    override fun getSourceName(): String {
        return "youtube"
    }

    override fun loadItem(manager: DefaultAudioPlayerManager, reference: AudioReference): AudioItem? {
        return try {
            loadItemOnce(reference)
        } catch (exception: FriendlyException) {
            // In case of a connection reset exception, try once more.
            if (HttpClientTools.isRetriableNetworkException(exception.cause)) {
                loadItemOnce(reference)
            } else {
                throw exception
            }
        }

    }

    private fun loadFromShortDomain(identifier: String): AudioItem {
        val urlInfo = getUrlInfo(identifier, true)
        return loadFromUrlWithVideoId(urlInfo.path.substring(1), urlInfo)
    }

    private fun loadFromMainDomain(identifier: String): AudioItem? {
        val urlInfo = getUrlInfo(identifier, true)
        println(urlInfo.parameters)
        if ("/watch" == urlInfo.path) {
            val videoId = urlInfo.parameters["v"]

            if (videoId != null) {
                return loadFromUrlWithVideoId(videoId, urlInfo)
            }
        } else if ("/playlist" == urlInfo.path) {
            val playlistId = urlInfo.parameters["list"]

            if (playlistId != null) {
                return loadPlaylistWithId(playlistId, null)
            }
        } else if ("/watch_videos" == urlInfo.path) {
            val videoIds = urlInfo.parameters["video_ids"]
            if (videoIds != null) {
                return loadAnonymous(videoIds)
            }
        }

        return null
    }

    override fun isTrackEncodable(track: AudioTrack): Boolean {
        return true
    }

    override fun encodeTrack(track: AudioTrack, output: DataOutput) {
        // No custom values that need saving
    }

    override fun decodeTrack(trackInfo: AudioTrackInfo, input: DataInput): AudioTrack {
        return YoutubeAudioTrack(trackInfo, this)
    }

    override fun shutdown() {
        ExceptionTools.closeWithWarnings(httpInterfaceManager)

        mixProvider.shutdown()
    }

    override fun configureRequests(configurator: java.util.function.Function<RequestConfig, RequestConfig>) {
        httpInterfaceManager.configureRequests(configurator)
        searchProvider.configureRequests(configurator)
    }

    override fun configureBuilder(configurator: Consumer<HttpClientBuilder>) {
        httpInterfaceManager.configureBuilder(configurator)
        searchProvider.configureBuilder(configurator)
    }

    private fun loadItemOnce(reference: AudioReference): AudioItem? {
        return loadNonSearch(reference.identifier) ?: searchProvider.loadSearchResult(
            reference.identifier.trim())

    }


    private fun loadAnonymous(videoIds: String): AudioItem? {
        try {
            httpInterface.use { httpInterface ->
                httpInterface.execute(HttpGet("https://www.youtube.com/watch_videos?video_ids=$videoIds"))
                    .use { response ->
                        val statusCode = response.statusLine.statusCode
                        val context = httpInterface.context
                        if (statusCode != 200) {
                            throw IOException("Invalid status code for playlist response: $statusCode")
                        }
                        // youtube currently transforms watch_video links into a link with a video id and a list id.
                        // because thats what happens, we can simply re-process with the redirected link
                        val redirects = context.redirectLocations
                        return if (redirects != null && !redirects.isEmpty()) {
                            loadNonSearch(redirects[0].toString())
                        } else {
                            throw FriendlyException(
                                "Unable to process youtube watch_videos link", FriendlyException.Severity.SUSPICIOUS,
                                IllegalStateException("Expected youtube to redirect watch_videos link to a watch?v={id}&list={list_id} link, but it did not redirect at all")
                            )
                        }
                    }
            }
        } catch (e: Exception) {
            throw ExceptionTools.wrapUnfriendlyExceptions(e)
        }

    }


    private fun loadFromUrlWithVideoId(videoId: String, urlInfo: UrlInfo): AudioItem {
        var videoId = videoId
        if (videoId.length > 11) {
            // YouTube allows extra junk in the end, it redirects to the correct video.
            videoId = videoId.substring(0, 11)
        }

        if (!directVideoIdPattern.matcher(videoId).matches()) {
            return AudioReference.NO_TRACK
        } else if (urlInfo.parameters.containsKey("list")) {
            val playlistId = urlInfo.parameters["list"]

            return if (playlistId!!.startsWith("RD")) {
                mixProvider.loadMixWithId(playlistId, videoId)
            } else {
                loadLinkedPlaylistWithId(urlInfo.parameters.getValue("list"), videoId)
            }
        } else {
            return loadTrackWithVideoId(videoId, false)
        }
    }

    private fun loadNonSearch(identifier: String): AudioItem? {
        println(identifier)
        for (extractor in extractors) {
            if (extractor.pattern.matcher(identifier).matches()) {
                val item = extractor.loader(identifier)

                if (item != null) {
                    return item
                }
            }
        }

        return null
    }

    private fun loadLinkedPlaylistWithId(playlistId: String, videoId: String): AudioItem {
        val playlist = loadPlaylistWithId(playlistId, videoId)

        return playlist ?: loadTrackWithVideoId(videoId, false)
    }

    private fun determineFailureReasonFromStatus(status: String?, reason: String?, mustExist: Boolean): Boolean {
        if ("fail" == status) {
            if (("This video does not exist." == reason || "This video is unavailable." == reason) && !mustExist) {
                return true
            } else if (reason != null) {
                throw FriendlyException(reason, FriendlyException.Severity.COMMON, null)
            }
        } else if ("ok" == status) {
            return false
        }

        throw FriendlyException(
            "Track is unavailable for an unknown reason.", FriendlyException.Severity.SUSPICIOUS,
            IllegalStateException("Main page had no video, but video info has no error.")
        )
    }

    @Throws(IOException::class)
    private fun getTrackInfoFromEmbedPage(httpInterface: HttpInterface, videoId: String): JsonBrowser {
        val basicInfo = loadTrackBaseInfoFromEmbedPage(httpInterface, videoId)
        basicInfo.put("args", loadTrackArgsFromVideoInfoPage(httpInterface, videoId, basicInfo.get("sts").text()))
        return basicInfo
    }

    @Throws(IOException::class)
    private fun loadTrackBaseInfoFromEmbedPage(httpInterface: HttpInterface, videoId: String): JsonBrowser {
        httpInterface.execute(HttpGet("https://www.youtube.com/embed/$videoId")).use { response ->
            val statusCode = response.statusLine.statusCode
            if (statusCode != 200) {
                throw IOException("Invalid status code for embed video page response: $statusCode")
            }

            val html = IOUtils.toString(response.entity.content, Charset.forName(CHARSET))
            val configJson = DataFormatTools.extractBetween(html, "'PLAYER_CONFIG': ", "});writeEmbed();")

            if (configJson != null) {
                return JsonBrowser.parse(configJson)
            }
        }

        throw FriendlyException(
            "Track information is unavailable.", FriendlyException.Severity.SUSPICIOUS,
            IllegalStateException("Expected player config is not present in embed page.")
        )
    }

    @Throws(IOException::class)
    private fun loadTrackArgsFromVideoInfoPage(
        httpInterface: HttpInterface,
        videoId: String,
        sts: String
    ): Map<String, String> {
        val videoApiUrl = "https://youtube.googleapis.com/v/$videoId"
        val encodedApiUrl = URLEncoder.encode(videoApiUrl, CHARSET)
        val url =
            "https://www.youtube.com/get_video_info?sts=" + sts + "&video_id=" + videoId + "&eurl=" + encodedApiUrl +
                    "hl=en_GB"

        httpInterface.execute(HttpGet(url)).use { response ->
            val statusCode = response.statusLine.statusCode
            if (statusCode != 200) {
                throw IOException("Invalid status code for video info response: $statusCode")
            }

            return convertToMapLayout(URLEncodedUtils.parse(response.entity))
        }
    }

    private fun loadPlaylistWithId(playlistId: String, selectedVideoId: String?): AudioPlaylist? {
        log.debug("Starting to load playlist with ID {}", playlistId)

        try {
            httpInterface.use { httpInterface ->
                httpInterface.execute(HttpGet("https://www.youtube.com/playlist?list=$playlistId")).use { response ->
                    val statusCode = response.statusLine.statusCode
                    if (statusCode != 200) {
                        throw IOException("Invalid status code for playlist response: $statusCode")
                    }

                    val document = Jsoup.parse(response.entity.content, CHARSET, "")
                    return buildPlaylist(httpInterface, document, selectedVideoId)
                }
            }
        } catch (e: Exception) {
            throw ExceptionTools.wrapUnfriendlyExceptions(e)
        }

    }

    @Throws(IOException::class)
    private fun buildPlaylist(
        httpInterface: HttpInterface,
        document: Document,
        selectedVideoId: String?
    ): AudioPlaylist? {
        val isAccessible = !document.select("#pl-header").isEmpty()

        if (!isAccessible) {
            return if (selectedVideoId != null) {
                null
            } else {
                throw FriendlyException("The playlist is private.", FriendlyException.Severity.COMMON, null)
            }
        }

        val container = document.select("#pl-header").first().parent()

        val playlistName = container.select(".pl-header-title").first().text()

        val tracks = ArrayList<AudioTrack>()
        var loadMoreUrl = extractPlaylistTracks(container, container, tracks)
        var loadCount = 0
        val pageCount = 6

        // Also load the next pages, each result gives us a JSON with separate values for list html and next page loader html
        while (loadMoreUrl != null && ++loadCount < pageCount) {
            httpInterface.execute(HttpGet("https://www.youtube.com$loadMoreUrl")).use { response ->
                val statusCode = response.statusLine.statusCode
                if (statusCode != 200) {
                    throw IOException("Invalid status code for playlist response: $statusCode")
                }

                val json = JsonBrowser.parse(response.entity.content)

                val html = json.get("content_html").text()
                val videoContainer = Jsoup.parse("<table>$html</table>", "")

                val moreHtml = json.get("load_more_widget_html").text()
                val moreContainer = if (moreHtml != null) Jsoup.parse(moreHtml) else null

                loadMoreUrl = extractPlaylistTracks(videoContainer, moreContainer, tracks)
            }
        }

        return BasicAudioPlaylist(playlistName, tracks, findSelectedTrack(tracks, selectedVideoId), false)
    }

    private fun extractPlaylistTracks(
        videoContainer: Element,
        loadMoreContainer: Element?,
        tracks: MutableList<AudioTrack>
    ): String? {
        for (video in videoContainer.select(".pl-video")) {
            val lengthElements = video.select(".timestamp span")

            // If the timestamp element does not exist, it means the video is private
            if (!lengthElements.isEmpty()) {
                val videoId = video.attr("data-video-id").trim { it <= ' ' }
                val title = video.attr("data-title").trim { it <= ' ' }
                val author = video.select(".pl-video-owner a").text().trim { it <= ' ' }
                val duration = DataFormatTools.durationTextToMillis(lengthElements.first().text())

                tracks.add(buildTrackObject(videoId, title, author, false, duration))
            }
        }

        if (loadMoreContainer != null) {
            val more = loadMoreContainer.select(".load-more-button")
            if (!more.isEmpty()) {
                return more.first().attr("data-uix-load-more-href")
            }
        }

        return null
    }


    private class UrlInfo constructor(val path: String, val parameters: Map<String, String>)

    private class Extractor constructor(
        val pattern: Pattern,
        val loader: (String) -> AudioItem?
    )

    companion object {
        private val log = LoggerFactory.getLogger(YoutubeAudioSourceManager::class.java)
        internal val CHARSET = "UTF-8"

        private const val PROTOCOL_REGEX = "(?:http://|https://|)"
        private const val DOMAIN_REGEX = "(?:www\\.|m\\.|music\\.|)youtube\\.com"
        private const val SHORT_DOMAIN_REGEX = "(?:www\\.|)youtu\\.be"
        private const val VIDEO_ID_REGEX = "(?<v>[a-zA-Z0-9_-]{11})"
        private const val PLAYLIST_ID_REGEX = "(?<list>(PL|LL|FL|UU)[a-zA-Z0-9_-]+)"

        private val directVideoIdPattern = Pattern.compile("^$VIDEO_ID_REGEX$")

        private fun getUrlInfo(url: String, retryValidPart: Boolean): UrlInfo {
            var url = url
            try {
                if (!url.startsWith("http://") && !url.startsWith("https://")) {
                    url = "https://$url"
                }

                val builder = URIBuilder(url)

                return UrlInfo(builder.path, builder.queryParams.map { it.name to it.value }.toMap())

            } catch (e: URISyntaxException) {
                return if (retryValidPart) {
                    getUrlInfo(url.substring(0, e.index - 1), false)
                } else {
                    throw FriendlyException("Not a valid URL: $url", FriendlyException.Severity.COMMON, e)
                }
            }

        }

    }
}