package main

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/schollz/progressbar/v3"
)

type M3u8Res struct {
	IsMaster bool
	NextUrl  string
	TsFile   []string
}

// 创建一个连接池
var httpClient = &http.Client{
	Transport: &http.Transport{
		MaxIdleConnsPerHost: 5,
	},
}

// 获取输入的m3u8链接
func getM3u8Path() (path string) {
	for {
		fmt.Println("请输入m3u8链接地址:")
		_, err := fmt.Scanln(&path)
		if err != nil {
			if err.Error() == "unexpected newline" {
				fmt.Println("你啥也没输入")
				continue
			}
			fmt.Println(err)
			continue
		}
		return
	}
}

// Get链接
func fetchM3u8Content(path string) string {
	res, err := httpClient.Get(path)
	if err != nil {
		fmt.Println(err)
		return ""
	}
	defer res.Body.Close()
	if res.StatusCode != 200 {
		fmt.Println("出了点问题", res.StatusCode)
		return ""
	}
	body, err := io.ReadAll(res.Body)
	if err != nil {
		fmt.Println(err)
		return ""
	}
	return string(body)
}

func parseM3u8(content string, baseUrl string) M3u8Res {
	lines := strings.Split(content, "\n")
	// 优先级最高：检查是否包含 Master 标签
	for i, line := range lines {
		if strings.Contains(line, "#EXT-X-STREAM-INF") {
			return parseMainContent(lines, i, baseUrl)
		}
	}
	return M3u8Res{
		IsMaster: false,
		TsFile:   parseSubContent(lines),
	}
}

func parseMainContent(lines []string, index int, baseUrl string) M3u8Res {
	// 取匹配标签的下一行作为 URL
	if index+1 < len(lines) {
		nextUrl := strings.TrimSpace(lines[index+1])
		return M3u8Res{
			IsMaster: true,
			NextUrl:  baseUrl + nextUrl,
		}
	}
	return M3u8Res{IsMaster: false}
}

func parseSubContent(lines []string) []string {
	var tsFiles []string
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if strings.HasPrefix(line, "#") || line == "" {
			continue
		}
		if strings.HasSuffix(line, ".jpeg") || strings.HasSuffix(line, ".ts") {
			tsFiles = append(tsFiles, line)
		}
	}
	return tsFiles
}

// 提取基础路径
func getBaseUrl(path string) string {
	lastSlash := strings.LastIndex(path, "/")
	if lastSlash == -1 {
		return path + "/"
	}
	return path[:lastSlash+1]
}

//func getNewBaseUrl(path string) string {
//	lastSlash := strings.LastIndex(path, "/")
//	if lastSlash == -1 {
//		return path + "/"
//	}
//	subPath := path[:lastSlash]
//	secondLastSlash := strings.LastIndex(subPath, "/")
//	return path[:secondLastSlash+1]
//}

// 拼接下载地址
func buildUrl(baseUrl string, m3u8TsFile []string) []string {
	var downLoadUrl []string
	for _, tsUrl := range m3u8TsFile {
		fullUrl := baseUrl + tsUrl
		downLoadUrl = append(downLoadUrl, fullUrl)
	}
	return downLoadUrl
}

// 协程用 index int 是用来标识每个文件的序号，在并发下载中特别重要
func downLoadOne(Url string, index int) error {
	maxRetry := 3 // 定义最大重试次数
	baseDelay := time.Second
	for retry := 0; retry < maxRetry; retry++ {
		if retry > 0 {
			waitTime := baseDelay * time.Duration(1<<uint(retry-1)) // 1s, 2s, 4s
			fmt.Printf("[%3d] 第 %d 次重试，等待 %v...\n", index, retry, waitTime)
			time.Sleep(waitTime)
		}
		resp, err := httpClient.Get(Url)
		if err != nil {
			fmt.Printf("连接失败: %v\n", err)
			if retry < maxRetry-1 {
				continue // 重试
			}
			return fmt.Errorf("（%d）连接失败 %d次后: %v", index, maxRetry, err)
		}
		if resp.StatusCode != 200 {
			resp.Body.Close()
			if retry < maxRetry-1 {
				continue // 重试
			}
			return fmt.Errorf("服务器状态异常: %d", resp.StatusCode)
		}
		fileName := fmt.Sprintf("%d", index)
		file, err := os.Create(fileName)
		if err != nil {
			resp.Body.Close()
			return err
		}

		_, err = io.Copy(file, resp.Body)
		closeErr := file.Close()
		resp.Body.Close()
		if err != nil || closeErr != nil {
			os.Remove(fileName)
			if err != nil {
				if retry < maxRetry-1 {
					continue // 重试
				}
				return fmt.Errorf("下载中断: %v %v", err, index)
			}
			return fmt.Errorf("文件关闭失败: %v", closeErr)
		}
		return nil
	}
	return nil
}

// 并行下载函数
func downLoadALL(fullUrl []string, maxConcurrent int) {
	var wait sync.WaitGroup
	total := len(fullUrl)
	bar := progressbar.Default(int64(total), "正在下载(%d)")
	maxT := make(chan struct{}, maxConcurrent) //定义一个长度为maxConcurrent的信道
	wait.Add(len(fullUrl))
	for i := range fullUrl {
		maxT <- struct{}{} // 写入一个信号
		go func(url string, index int) {
			defer func() {
				<-maxT
				bar.Add(1)
			}()
			defer wait.Done()
			err := downLoadOne(url, index)
			if err != nil {
				fmt.Printf("文件 %d 失败: %v ", index, err)
			}
		}(fullUrl[i], i)
	}
	wait.Wait()
	bar.Finish()
}

// 碎片拼接
func videoBuild(num int, outputpath string, deleteAfter bool) error {
	output, err := os.Create(outputpath)
	if err != nil {
		return err
	}
	defer output.Close()

	for i := 0; i < num; i++ {
		fileName := fmt.Sprintf("%d", i)
		input, err := os.Open(fileName)
		if err != nil {
			return fmt.Errorf("打开文件失败\n")
		}
		writtenBytes, err := io.Copy(output, input)
		input.Close()
		if err != nil {
			return fmt.Errorf("合并文件失败\n")
		}
		if deleteAfter {
			os.Remove(fileName)
		}
		fmt.Printf("合并成功,%d\n", writtenBytes)
	}

	return nil
}

func main() {
	m3u8Path := getM3u8Path()              // 输入链接
	m3u8Body := fetchM3u8Content(m3u8Path) // get到链接文件
	m3u8BaseUrl := getBaseUrl(m3u8Path)    // 获取基础路径
	m3u8TsFile := parseM3u8(m3u8Body, m3u8BaseUrl)
	if m3u8TsFile.IsMaster {
		fmt.Println("检测到主播放列表，准备跳转:", m3u8TsFile.NextUrl)
		baseNew := getBaseUrl(m3u8TsFile.NextUrl)       // 新的基础路径
		m3u8Body = fetchM3u8Content(m3u8TsFile.NextUrl) // get新的内容
		m3u8TsFileNew := parseM3u8(m3u8Body, baseNew)
		fullUrl := buildUrl(baseNew, m3u8TsFileNew.TsFile)
		downLoadALL(fullUrl, 10)
		fmt.Printf("下载完成？开始合并...\n")
		videoBuild(len(m3u8TsFileNew.TsFile), "test.mp4", true)
	} else {
		fmt.Println("获得分片列表，准备下载:", len(m3u8TsFile.TsFile), "个文件")
		fullUrl := buildUrl(m3u8BaseUrl, m3u8TsFile.TsFile)
		downLoadALL(fullUrl, 10)
		fmt.Printf("下载完成？开始合并...\n")
		videoBuild(len(m3u8TsFile.TsFile), "test.mp4", true)
	}
}

//http://devimages.apple.com/iphone/samples/bipbop/gear1/prog_index.m3u8
//https://test-streams.mux.dev/x36xhzz/x36xhzz.m3u8
//https://test-streams.mux.dev/x36xhzz/url_0/193039199_mp4_h264_aac_hd_7.m3u8
//https://surrit.com/0831edf7-98aa-49ff-a865-42777726e282/720p/video.m3u8