func main() {
	fileURL := flag.String("u", "", "downloade url of the file")
	flag.Parse()
	if *fileURL == "" {
		log.Println("Please input a download url")
		flag.Usage()
		return
	}
	fileDir, err := os.Getwd()
	if err != nil {
		log.Println(err)
		return
	}
	// 下載文件保存路徑
	filePath := filepath.Join(fileDir, filepath.Base(*fileURL))
	err = downloadFile(*fileURL, filePath)
	if err != nil {
		log.Println(err)
		return
	}
	log.Println("download file success:", filePath)
}

// 下載文件
func downloadFile(fileURL string, filePath string) error {
	log.Println("downloading file:", fileURL, "to", filePath)
	taskCh := make(chan [2]int64, runtime.NumCPU())
	wg := new(sync.WaitGroup)
	// 創(chuàng)建執(zhí)行下載任務的 worker
	err := initWorker(fileURL, filePath, taskCh, wg)
	if err != nil {
		return fmt.Errorf("init worker failed: %v", err)
	}
	// 分發(fā)下載任務
	err = dispatchTask(fileURL, taskCh)
	if err != nil {
		return fmt.Errorf("dispacth task failed: %v", err)
	}
	// 等待所有下載任務完成
	wg.Wait()
	return nil
}

// 初始化 下載 worker
func initWorker(url string, filePath string, taskCh chan [2]int64, wg *sync.WaitGroup) error {
	for i := 0; i < runtime.NumCPU(); i++ {
		// 打開文件句柄
		file, err := os.OpenFile(filePath, os.O_CREATE|os.O_RDWR, 0644)
		if err != nil {
			return err
		}
		wg.Add(1)
		go func(file *os.File, taskCh chan [2]int64) {
			defer wg.Done()
			defer file.Close()
			// 循環(huán)從 taskCh 中獲取下載任務并下載
			for part := range taskCh {
				log.Printf("downloading part, start offset: %d, end offset: %d", part[0], part[1])
				// 重試下載，最大重試次數(shù)為 10 次，每次下載失敗后等待 1 秒
				err := retryWithWaitTime(10, func() error {
					return downloadPart(url, file, part[0], part[1])
				}, time.Second)
				if err != nil {
					log.Printf("download part %d failed: %v", part, err)
				}
			}
		}(file, taskCh)
	}
	return nil
}

// 分發(fā)下載任務
func dispatchTask(url string, taskCh chan [2]int64) error {
	defer close(taskCh)
	fileSize, err := getFileSize(url)
	if err != nil {
		return err
	}
	// 分片大小 1MB
	const chunkSize = 1024 * 1024
	parts := fileSize / chunkSize
	log.Println("file size:", fileSize, "parts:", parts, "chunk size:", chunkSize)
	for i := int64(0); i < parts; i++ {
		// 計算分片的起始和結束位置
		startOffset := i * chunkSize
		endOffset := startOffset + chunkSize - 1
		// 發(fā)送下載任務
		taskCh <- [2]int64{startOffset, endOffset}
	}
	// 發(fā)送最后一個分片的下載任務
	if fileSize % chunkSize != 0 {
		taskCh <- [2]int64{parts * chunkSize, fileSize - 1}
	}
	return nil
}

// 獲取文件大小
func getFileSize(url string) (int64, error) {
	resp, err := http.Head(url)
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()
	return resp.ContentLength, nil
}

// 下載文件分片
func downloadPart(url string, file *os.File, startPos, endPos int64) error {
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return err
	}
	// 設置文件分片區(qū)間的請求頭
	req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", startPos, endPos))
	resp, err := http.DefaultTransport.RoundTrip(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	// 如果服務器返回的狀態(tài)碼不是 206 Partial Content，則說明下載失敗
	if resp.StatusCode != http.StatusPartialContent {
		data, err := io.ReadAll(resp.Body)
		if err != nil {
			return err
		}
		log.Println("unexpected data:", string(data))
		return fmt.Errorf("unexpected status code: %d", resp.StatusCode)
	}
	// 文件指針移動到分片的起始位置
	_, err = file.Seek(startPos, 0)
	if err != nil {
		return err
	}
	// 寫入分片數(shù)據(jù)到文件
	_, err = io.Copy(file, resp.Body)
	if err != nil {
		return err
	}
	return nil
}

// 重試函數(shù)
func retryWithWaitTime(retryCount int, fn func() error, waitTime time.Duration) error {
	var err error
	for i := 0; i < retryCount; i++ {
		e := fn()
		if e != nil {
			errors.Join(err, e)
			time.Sleep(waitTime)
			continue
		}
		return nil
	}
	return err
}