package dylive

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"
)

const userAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:109.0) Gecko/20100101 Firefox/115.0"

type (
	Category struct {
		Id         string
		Name       string
		Categories []Category
	}

	dyCategory struct {
		Partition struct {
			IDStr string `json:"id_str"`
			Type  int    `json:"type"`
			Title string `json:"title"`
		} `json:"partition"`
		SubPartition []dyCategory `json:"sub_partition"`
	}

	dyliveCategories struct {
		CategoryData []dyCategory `json:"categoryData"`
	}
)

// GetCategories gets all Douyin live stream categories.
func GetCategories(ctx context.Context) ([]Category, error) {
	const first = "4_101"
	var categories []Category
	data, err := getCategoryPageData(ctx, first, "categoryData")
	if err != nil {
		return nil, err
	}
	var cats dyliveCategories
	if err := getDataInArray(data[0], &cats); err != nil {
		return nil, err
	}
	categories = append(categories, deepConvertDyCategories(cats.CategoryData, nil)...)
	return categories, nil
}

func deepConvertDyCategories(in []dyCategory, whitelist []string) []Category {
	out := []Category{}
	for _, c := range in {
		if cat := convertDyCategory(c, nil, 0, whitelist); cat != nil {
			out = append(out, *cat)
		}
	}
	return out
}

func convertDyCategory(c dyCategory, parentIDParts []string, level int, whitelist []string) *Category {
	id := fmt.Sprintf("%d_%s", c.Partition.Type, c.Partition.IDStr)
	if len(whitelist) > 0 && (level >= len(whitelist) || whitelist[level] != id) {
		return nil
	}
	idParts := append([]string{}, parentIDParts...)
	idParts = append(idParts, id)
	fullId := strings.Join(idParts, "_")
	children := []Category{}
	for _, sub := range c.SubPartition {
		if cat := convertDyCategory(sub, idParts, level+1, whitelist); cat != nil {
			children = append(children, *cat)
		}
	}
	return &Category{
		Id:         fullId,
		Name:       c.Partition.Title,
		Categories: children,
	}
}

const (
	RoomStatusLiveOn RoomStatus = 2 + iota
	_
	RoomStatusLiveOff
)

type (
	RoomStatus = int

	Room struct {
		Id                string
		DouyinId          string
		StatusCode        RoomStatus
		Name              string
		CoverUrl          string
		WebUrl            string
		CurrentUsersCount string
		TotalUsersCount   string
		Category          *Category
		User              User
		StreamUrl         string
		FlvStreamUrls     map[string]string
		HlsStreamUrls     map[string]string
	}

	User struct {
		Name    string
		Picture string
	}

	dyUser struct {
		Nickname    string `json:"nickname"`
		AvatarThumb struct {
			UrlList []string `json:"url_list"`
		} `json:"avatar_thumb"`
	}

	dyliveRoom struct {
		IdStr  string `json:"id_str"`
		Title  string `json:"title"`
		Status int    `json:"status"`
		Cover  struct {
			UrlList []string `json:"url_list"`
		} `json:"cover"`
		Stats struct {
			TotalUserStr string `json:"total_user_str"`
			UserCountStr string `json:"user_count_str"`
		} `json:"stats"`
		Owner     dyUser `json:"owner"`
		StreamUrl struct {
			FlvPullUrl        map[string]string `json:"flv_pull_url"`
			HlsPullUrlMap     map[string]string `json:"hls_pull_url_map"`
			DefaultResolution string            `json:"default_resolution"`
		} `json:"stream_url"`
		RoomViewStats struct {
			DisplayValue int `json:"display_value"`
		} `json:"room_view_stats"`
	}

	dyliveCategory struct {
		RoomsData struct {
			Data []struct {
				Room      dyliveRoom `json:"room"`
				WebRid    string     `json:"web_rid"`
				StreamSrc string     `json:"streamSrc"`
				Cover     string     `json:"cover"`
				Avatar    string     `json:"avatar"`
			} `json:"data"`
		} `json:"roomsData"`
		CategoryData []dyCategory `json:"categoryData"`
		CategoryList []string     `json:"categoryList"`
	}
)

// FlvUrlForQuality returns the .flv stream URL for the given quality (uhd, hd, ld, sd).
// If no matching URL is found, it returns the room's default StreamUrl.
func (room Room) FlvUrlForQuality(quality string) string {
	return room.urlForQuality(room.FlvStreamUrls, quality)
}

// HlsUrlForQuality returns the .m3u8 stream URL for the given quality (uhd, hd, ld, sd).
// If no matching URL is found, it returns the room's default StreamUrl.
func (room Room) HlsUrlForQuality(quality string) string {
	return room.urlForQuality(room.HlsStreamUrls, quality)
}

func (room Room) urlForQuality(urls map[string]string, quality string) string {
	quality = strings.ToLower(quality)
	for key, value := range urls {
		switch quality {
		case "uhd":
			if strings.Contains(key, "FULL_HD") || strings.Contains(value, "_uhd") {
				return value
			}
		case "hd":
			if strings.Contains(value, "_hd") {
				return value
			}
		case "ld":
			if strings.Contains(value, "_ld") {
				return value
			}
		case "sd":
			if strings.Contains(value, "_sd") {
				return value
			}
		default:
			return room.StreamUrl
		}
	}
	return room.StreamUrl
}

// GetRoomsByCategory gets top 15 Douyin live stream rooms of a category.
func GetRoomsByCategory(ctx context.Context, categoryId string) ([]Room, error) {
	data, err := getCategoryPageData(ctx, categoryId, "roomsData")
	if err != nil {
		return nil, err
	}
	roomsData := data[0]
	var cat dyliveCategory
	if err := getDataInArray(roomsData, &cat); err != nil {
		return nil, err
	}
	var category *Category
	if categories := deepConvertDyCategories(cat.CategoryData, cat.CategoryList); len(categories) > 0 {
		category = &categories[0]
	}

	var rooms []Room
	for _, room := range cat.RoomsData.Data {
		var count string
		if room.Room.RoomViewStats.DisplayValue > 0 {
			count = strconv.Itoa(room.Room.RoomViewStats.DisplayValue)
		} else {
			count = room.Room.Stats.UserCountStr
		}
		rooms = append(rooms, Room{
			Id:                room.Room.IdStr,
			DouyinId:          room.WebRid,
			StatusCode:        RoomStatusLiveOn,
			Name:              room.Room.Title,
			CoverUrl:          room.Cover,
			WebUrl:            "https://live.douyin.com/" + room.WebRid,
			StreamUrl:         room.StreamSrc,
			FlvStreamUrls:     room.Room.StreamUrl.FlvPullUrl,
			HlsStreamUrls:     room.Room.StreamUrl.HlsPullUrlMap,
			CurrentUsersCount: count,
			TotalUsersCount:   room.Room.Stats.TotalUserStr,
			Category:          category,
			User: User{
				Name:    room.Room.Owner.Nickname,
				Picture: room.Avatar,
			},
		})
	}
	return rooms, nil
}

func getCategoryPageData(ctx context.Context, id string, filters ...string) ([]string, error) {
	req, err := http.NewRequestWithContext(ctx, "GET", "https://live.douyin.com/categorynew/"+id, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", userAgent)
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	parts := getDataInHtml(string(b))
	var output []string
	for _, filter := range filters {
		var ret string
		for _, part := range parts {
			if strings.Contains(part, filter) {
				ret = part
				break
			}
		}
		output = append(output, ret)
	}
	return output, nil
}

type (
	dyliveRoomDetails struct {
		State struct {
			RoomStore struct {
				RoomInfo struct {
					Room   dyliveRoom `json:"room"`
					WebRid string     `json:"web_rid"`
					Anchor dyUser     `json:"anchor"`
				} `json:"roomInfo"`
			} `json:"roomStore"`
		} `json:"state"`
	}
)

// GetRoom get live stream room details by Douyin ID (抖音号)
func GetRoom(ctx context.Context, douyinId string) (*Room, error) {
	data, err := getLivePageData(ctx, douyinId, "flv_pull_url")
	if err != nil {
		return nil, err
	}
	if len(data) == 0 || data[0] == "" {
		return nil, fmt.Errorf("DouyinId %s does not exist", douyinId)
	}
	roomsData := data[0]
	var page dyliveRoomDetails
	if err := getDataInArray(roomsData, &page); err != nil {
		return nil, err
	}

	info := page.State.RoomStore.RoomInfo

	var cover string
	if len(info.Room.Cover.UrlList) > 0 {
		cover = info.Room.Cover.UrlList[0]
	}

	streamUrl := info.Room.StreamUrl.FlvPullUrl[info.Room.StreamUrl.DefaultResolution]

	var count string
	if info.Room.RoomViewStats.DisplayValue > 0 {
		count = strconv.Itoa(info.Room.RoomViewStats.DisplayValue)
	} else {
		count = info.Room.Stats.UserCountStr
	}

	userName := info.Room.Owner.Nickname
	if userName == "" {
		userName = info.Anchor.Nickname
	}

	var userPicture string
	if len(info.Room.Owner.AvatarThumb.UrlList) > 0 {
		userPicture = info.Room.Owner.AvatarThumb.UrlList[0]
	} else if len(info.Anchor.AvatarThumb.UrlList) > 0 {
		userPicture = info.Anchor.AvatarThumb.UrlList[0]
	}

	return &Room{
		Id:                info.Room.IdStr,
		DouyinId:          info.WebRid,
		StatusCode:        info.Room.Status,
		Name:              info.Room.Title,
		CoverUrl:          cover,
		WebUrl:            "https://live.douyin.com/" + info.WebRid,
		StreamUrl:         streamUrl,
		FlvStreamUrls:     info.Room.StreamUrl.FlvPullUrl,
		HlsStreamUrls:     info.Room.StreamUrl.HlsPullUrlMap,
		CurrentUsersCount: count,
		TotalUsersCount:   info.Room.Stats.TotalUserStr,
		User: User{
			Name:    userName,
			Picture: userPicture,
		},
	}, nil
}

func getLivePageData(ctx context.Context, douyinId string, filters ...string) ([]string, error) {
	req, err := http.NewRequestWithContext(ctx, "GET", "https://live.douyin.com/"+douyinId, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", userAgent)
	req.Header.Set("Cookie", "__ac_nonce=064caded4009deafd8b89")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}
	parts := getDataInHtml(string(b))
	var output []string
	for _, filter := range filters {
		var ret string
		for _, part := range parts {
			if strings.Contains(part, filter) {
				ret = part
				break
			}
		}
		output = append(output, ret)
	}
	return output, nil
}

func getDataInHtml(input string) (output []string) {
	const funcName = "__pace_f"
	const endTag = "</script>"
	var parts []string
	for {
		a := strings.Index(input, funcName)
		if a == -1 {
			break
		}
		input = input[a+len(funcName):]
		b := strings.Index(input, `"`)
		if b < 0 {
			continue
		}
		input = input[b+1:]
		b = strings.Index(input, endTag)
		if b < 0 {
			continue
		}
		b = strings.LastIndex(input[:b], `"`)
		if b < 0 {
			continue
		}
		var ret string
		if json.Unmarshal([]byte(`"`+input[:b]+`"`), &ret) != nil {
			continue
		}
		parts = append(parts, ret)
	}
	parts = strings.Split(strings.Join(parts, "\n"), "\n")
	for _, part := range parts {
		a := strings.IndexAny(part, "[{")
		if a == -1 {
			continue
		}
		b := strings.LastIndexAny(part, "}]")
		if b == -1 {
			continue
		}
		output = append(output, part[a:b+1])
	}
	return
}

func getDataInArray(input string, target interface{}) error {
	var array []interface{}
	if err := json.Unmarshal([]byte(input), &array); err != nil {
		return err
	}
	for _, element := range array {
		switch v := element.(type) {
		case map[string]interface{}:
			jsonStr, err := json.Marshal(v)
			if err != nil {
				continue
			}
			return json.Unmarshal(jsonStr, target)
		}
	}
	return nil
}
