/* SPDX-License-Identifier: MIT
 *
 * Copyright (C) 2020 jet tsang zeon-git. All Rights Reserved.
 */

package main

import (
	"bufio"
	"flag"
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/antchfx/htmlquery"
	"github.com/valyala/fasthttp"
)

type ShortMap map[string]string

var listUrl string
var blockTxt string
var shortenMap map[string]string
var listen string
var reloadURI = "/reload"
var clientHeaderURI = "/header"
var UserAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/103.0.5060.53 Safari/537.36"

func init() {
	const (
		urlUsage     = "This program will get a list from the url you provide."
		blockUsage   = "html block to query."
		defaultBlock = ""

		defaultListen = "localhost:9530"
		listenUsage   = "TCP address to listen to"
	)
	flag.StringVar(&listUrl, "url", "https://pastebin.com/raw/9g4msZxC", urlUsage)
	flag.StringVar(&listUrl, "u", "https://pastebin.com/raw/9g4msZxC", urlUsage+" (shorthand)")

	flag.StringVar(&blockTxt, "block", defaultBlock, blockUsage)
	flag.StringVar(&blockTxt, "b", defaultBlock, blockUsage+" (shorthand)")

	flag.StringVar(&listen, "listen", defaultListen, listenUsage)
	flag.StringVar(&listen, "l", defaultListen, listenUsage+" (shorthand)")
	shortenMap = map[string]string{}
}
func main() {
	flag.Parse()
	go autoLoad()
	log.Println("Listening to:", listen)
	s := &fasthttp.Server{
		Handler:               requestHandler,
		NoDefaultServerHeader: true,
	}

	if err := s.ListenAndServe(listen); err != nil {
		log.Fatalf("Error in ListenAndServe: %s", err)
	}
}

func requestHandler(ctx *fasthttp.RequestCtx) {
	strr := string(ctx.RequestURI())
	ctx.Response.Header.Add("Server", "jet's relay server")
	if strr == reloadURI {
		loadData()
		fmt.Fprint(ctx, "ok")
		ctx.SetStatusCode(fasthttp.StatusAccepted)
		return
	} else if strr == clientHeaderURI {
		fmt.Fprint(ctx, string(ctx.Request.Header.Header()))
		ctx.SetStatusCode(fasthttp.StatusAccepted)
		return
	} else if strings.HasPrefix(strr, clientHeaderURI+"/") {
		headerkey := strings.TrimLeft(strr, clientHeaderURI+"/")
		log.Println("headerkey", headerkey)
		fmt.Fprint(ctx, string(ctx.Request.Header.Peek(headerkey)))
		ctx.SetStatusCode(fasthttp.StatusAccepted)
		return
	} else if strings.HasPrefix(strr, "/yt/") {
		var req fasthttp.Request
		var resp fasthttp.Response
		ythash := strings.TrimLeft(strr, "/yt/")
		req.Header.Add("User-Agent", UserAgent)
		req.Header.Add("Referer", "https://www.youtube.com")
		req.SetRequestURI("https://www.youtube.com/watch?v=" + ythash)
		fasthttp.Do(&req, &resp)
		ytbody := string(resp.Body())
		m3u8 := GetBetweenStr(ytbody, "hlsManifestUrl\":\"", "\"},\"heart")
		ctx.Redirect(m3u8, 302)
		defer func() {
			fasthttp.ReleaseRequest(&req)
			req.ConnectionClose()
			resp.ConnectionClose()
		}()
		return
	}
	if relayUrl, exist := shortenMap[string(ctx.Path())]; exist {
		var req fasthttp.Request
		ctx.Request.Header.CopyTo(&req.Header)
		req.SetRequestURI(relayUrl)
		req.Header.SetBytesV("Host", req.Host())
		ctx.URI().QueryArgs().VisitAll(func(key []byte, value []byte) {
			req.URI().QueryArgs().AddBytesKV(key, value)
		})
		defer fasthttp.ReleaseRequest(&req)

		fasthttp.Do(&req, &ctx.Response)
	} else {
		ctx.Error("err", 404)
	}
}

func autoLoad() {
	for {
		loadData()
		time.Sleep(time.Hour)
	}
}

func loadData() {
	log.Println("Loading data.", listUrl)
	var code string

	if blockTxt != "" {
		doc, err := htmlquery.LoadURL(listUrl)
		if err != nil {
			log.Println("url err:", listUrl)
			return
		}
		found, err := htmlquery.QueryAll(doc, blockTxt)
		if err != nil {
			log.Println("block not found", err)
			return
		}
		if len(found) == 0 {
			log.Println("No elements found")
			return
		}
		code = htmlquery.InnerText(found[0])
	} else {
		var dst []byte
		statusCode, byteCode, err := fasthttp.Get(dst, string(listUrl))
		if err != nil || statusCode != fasthttp.StatusOK {
			log.Printf("HTTP request failed with status code %d\n", statusCode)
			return
		}
		code = string(byteCode)
	}

	scanner := bufio.NewScanner(strings.NewReader(code))
	for scanner.Scan() {
		maplist := strings.Split(scanner.Text(), "#")
		if len(maplist) < 2 {
			log.Println("Invalid data format:", scanner.Text())
			continue // Skip this iteration as the data format is not valid
		}
		shortenMap[maplist[0]] = maplist[1]
	}
	log.Println("URLs mapped.")
}

func GetBetweenStr(str, start, end string) string {
	n := strings.Index(str, start) + len(start)
	if n == -1 {
		n = 0
	}
	str = string([]byte(str)[n:])
	m := strings.Index(str, end)
	if m == -1 {
		m = len(str)
	}
	str = string([]byte(str)[:m])
	return str
}