package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/glebarez/sqlite"
	"github.com/studio-b12/gowebdav"
	"gorm.io/gorm"
)

type Task struct {
	ID              uint      `gorm:"primaryKey" json:"id"`
	Source          string    `gorm:"uniqueIndex:idx_tasks_source_rid_quality" json:"source"`
	RID             string    `gorm:"uniqueIndex:idx_tasks_source_rid_quality" json:"rid"`
	Title           string    `json:"title"`
	Artist          string    `json:"artist"`
	Album           string    `json:"album"`
	Pic             string    `json:"pic"`
	Quality         string    `gorm:"uniqueIndex:idx_tasks_source_rid_quality" json:"quality"`
	ResolvedQuality string    `json:"resolvedQuality,omitempty"`
	ActualQuality   string    `json:"actualQuality,omitempty"`
	Format          string    `json:"format"`
	Status          string    `json:"status"`
	Error           string    `json:"error,omitempty"`
	LocalPath       string    `json:"localPath,omitempty"`
	RemotePath      string    `json:"remotePath,omitempty"`
	TotalBytes      int64     `json:"totalBytes"`
	DownloadedBytes int64     `json:"downloadedBytes"`
	Progress        float64   `json:"progress"`
	CreatedAt       time.Time `json:"createdAt"`
	UpdatedAt       time.Time `json:"updatedAt"`
}

type Setting struct {
	ID    uint   `gorm:"primaryKey"`
	Key   string `gorm:"uniqueIndex"`
	Value string
}
type SongRequest struct {
	Source  string `json:"source"`
	RID     string `json:"rid"`
	Title   string `json:"title"`
	Artist  string `json:"artist"`
	Album   string `json:"album"`
	Pic     string `json:"pic"`
	Quality string `json:"quality"`
}
type Config struct {
	AMAPI, DownloadDir, WebDAVURL, WebDAVUser, WebDAVPass, WebDAVPath, DefaultQuality string
	WorkerCount                                                                       int
}
type App struct {
	db         *gorm.DB
	cfg        Config
	queue      chan uint
	wg         sync.WaitGroup
	httpClient *http.Client
}

type playInfo struct {
	URL     string            `json:"url"`
	Level   string            `json:"level"`
	Format  string            `json:"format"`
	Headers map[string]string `json:"headers"`
}

func main() {
	dbPath := getenv("DB_PATH", "/data/downloader.db")
	_ = os.MkdirAll(filepath.Dir(dbPath), 0755)
	db, err := gorm.Open(sqlite.Open(dbPath), &gorm.Config{})
	must(err)
	must(db.AutoMigrate(&Task{}, &Setting{}))
	a := &App{db: db, cfg: Config{AMAPI: getenv("AM_API_URL", "http://127.0.0.1:9530/api"), DownloadDir: getenv("DOWNLOAD_DIR", "/downloads"), DefaultQuality: getenv("DEFAULT_QUALITY", "standard"), WorkerCount: atoi(getenv("WORKERS", "2"), 2)}, queue: make(chan uint, 100), httpClient: &http.Client{Timeout: 2 * time.Minute}}
	_ = os.MkdirAll(a.cfg.DownloadDir, 0755)
	a.loadSettings()
	a.startWorkers()
	r := gin.Default()
	r.StaticFile("/", "/app/index.html")
	api := r.Group("/api")
	api.GET("/search", a.search)
	api.GET("/tasks", a.listTasks)
	api.POST("/tasks", a.addTask)
	api.POST("/tasks/:id/retry", a.retryTask)
	api.DELETE("/tasks/:id", a.deleteTask)
	api.GET("/settings", a.getSettings)
	api.PUT("/settings", a.putSettings)
	api.POST("/webdav/test", a.testWebDAV)
	must(r.Run(":" + getenv("PORT", "9540")))
}

func (a *App) loadSettings() {
	var ss []Setting
	a.db.Find(&ss)
	for _, s := range ss {
		switch s.Key {
		case "am_api":
			a.cfg.AMAPI = s.Value
		case "download_dir":
			a.cfg.DownloadDir = s.Value
		case "webdav_url":
			a.cfg.WebDAVURL = s.Value
		case "webdav_user":
			a.cfg.WebDAVUser = s.Value
		case "webdav_pass":
			a.cfg.WebDAVPass = s.Value
		case "webdav_path":
			a.cfg.WebDAVPath = s.Value
		case "quality":
			a.cfg.DefaultQuality = s.Value
		}
	}
}
func (a *App) save(k, v string) {
	var s Setting
	if a.db.Where("key = ?", k).First(&s).Error != nil {
		a.db.Create(&Setting{Key: k, Value: v})
	} else {
		s.Value = v
		a.db.Save(&s)
	}
}
func (a *App) startWorkers() {
	for i := 0; i < a.cfg.WorkerCount; i++ {
		a.wg.Add(1)
		go func() {
			defer a.wg.Done()
			for id := range a.queue {
				a.download(id)
			}
		}()
	}
}
func (a *App) addTask(c *gin.Context) {
	var in SongRequest
	if c.BindJSON(&in) != nil || in.RID == "" || in.Title == "" {
		c.JSON(400, gin.H{"error": "source/rid/title required"})
		return
	}
	if in.Source == "" {
		in.Source = "kw"
	}
	if in.Quality == "" {
		in.Quality = a.cfg.DefaultQuality
	}
	var existing Task
	err := a.db.Where("source = ? AND r_id = ? AND quality = ?", in.Source, in.RID, in.Quality).Order("id desc").First(&existing).Error
	if err == nil {
		if existing.Status != "error" {
			c.JSON(200, existing)
			return
		}
		a.db.Model(&existing).Updates(map[string]any{"title": in.Title, "artist": in.Artist, "album": in.Album, "pic": in.Pic, "status": "pending", "error": ""})
		a.queue <- existing.ID
		c.JSON(202, existing)
		return
	}
	t := Task{Source: in.Source, RID: in.RID, Title: in.Title, Artist: in.Artist, Album: in.Album, Pic: in.Pic, Quality: in.Quality, Status: "pending"}
	if err := a.db.Create(&t).Error; err != nil {
		// A concurrent request may have inserted the same task; return that task.
		if a.db.Where("source = ? AND r_id = ? AND quality = ?", in.Source, in.RID, in.Quality).Order("id desc").First(&existing).Error == nil {
			c.JSON(200, existing)
			return
		}
		c.JSON(500, gin.H{"error": err.Error()})
		return
	}
	a.queue <- t.ID
	c.JSON(202, t)
}
func (a *App) search(c *gin.Context) {
	keyword := strings.TrimSpace(c.Query("keyword"))
	source := c.DefaultQuery("source", "kw")
	if keyword == "" {
		c.JSON(400, gin.H{"error": "keyword required"})
		return
	}
	u, _ := url.Parse(strings.TrimRight(a.cfg.AMAPI, "/") + "/search")
	q := u.Query()
	q.Set("keyword", keyword)
	q.Set("source", source)
	q.Set("page", c.DefaultQuery("page", "1"))
	q.Set("limit", c.DefaultQuery("limit", "20"))
	u.RawQuery = q.Encode()
	resp, err := a.httpClient.Get(u.String())
	if err != nil {
		c.JSON(502, gin.H{"error": err.Error()})
		return
	}
	defer resp.Body.Close()
	c.Status(resp.StatusCode)
	_, _ = io.Copy(c.Writer, resp.Body)
}

func (a *App) listTasks(c *gin.Context) {
	var ts []Task
	a.db.Order("id desc").Limit(100).Find(&ts)
	c.JSON(200, gin.H{"data": ts})
}
func (a *App) retryTask(c *gin.Context) {
	id, _ := strconv.ParseUint(c.Param("id"), 10, 64)
	var t Task
	if a.db.First(&t, id).Error != nil {
		c.Status(404)
		return
	}
	a.db.Model(&t).Updates(map[string]any{"status": "pending", "error": "", "progress": 0, "downloaded_bytes": 0})
	t.Status = "pending"
	t.Error = ""
	a.queue <- t.ID
	c.JSON(202, t)
}
func (a *App) deleteTask(c *gin.Context) {
	id, _ := strconv.ParseUint(c.Param("id"), 10, 64)
	a.db.Delete(&Task{}, id)
	c.Status(204)
}

func (a *App) download(id uint) {
	var t Task
	if a.db.First(&t, id).Error != nil {
		return
	}
	// A task may be queued twice; only pending tasks are eligible to start.
	if t.Status != "pending" {
		return
	}
	a.db.Model(&t).Updates(map[string]any{"status": "resolving", "error": "", "progress": 0})
	info, err := a.resolve(t)
	if err != nil {
		a.fail(&t, err)
		return
	}
	priorLocal := t.LocalPath
	ext := normalizedFormat(info.Format)
	if ext == "" {
		ext = normalizedFormat(filepath.Ext(info.URL))
	}
	if ext == "" {
		ext = "mp3"
	}
	filename := safe(t.Artist) + " - " + safe(t.Title) + "." + ext
	local := filepath.Join(a.cfg.DownloadDir, filename)
	t.Format = ext
	t.ResolvedQuality = info.Level
	t.ActualQuality = ""
	t.LocalPath = local
	if priorLocal != "" {
		if fi, statErr := os.Stat(priorLocal); statErr == nil && fi.Size() > 0 {
			local = priorLocal
			t.LocalPath = local
		}
	}
	if fi, statErr := os.Stat(local); statErr == nil && fi.Size() > 0 {
		actual, detectErr := detectAudioFormat(local)
		if detectErr == nil {
			actual = normalizedFormat(actual)
			if actual != t.Format {
				newLocal, renameErr := renameToActualFormat(local, actual)
				if renameErr != nil {
					a.fail(&t, renameErr)
					return
				}
				local = newLocal
				t.LocalPath = local
			}
			t.Format = actual
			t.ActualQuality = actualQuality(t.Quality, t.ResolvedQuality, actual)
		}
		t.TotalBytes = fi.Size()
		t.DownloadedBytes = fi.Size()
		t.Progress = 100
		a.db.Save(&t)
		a.upload(&t)
		return
	}
	a.db.Model(&t).Updates(map[string]any{"status": "downloading", "format": ext, "local_path": local})
	req, err := http.NewRequest("GET", info.URL, nil)
	if err != nil {
		a.fail(&t, err)
		return
	}
	for k, v := range info.Headers {
		req.Header.Set(k, v)
	}
	resp, err := a.httpClient.Do(req)
	if err != nil {
		a.fail(&t, err)
		return
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		a.fail(&t, fmt.Errorf("media http %d", resp.StatusCode))
		return
	}
	t.Status = "downloading"
	t.TotalBytes = resp.ContentLength
	t.DownloadedBytes = 0
	t.Progress = 0
	a.db.Save(&t)
	tmp := local + ".part"
	f, err := os.Create(tmp)
	if err != nil {
		a.fail(&t, err)
		return
	}
	var n int64
	buf := make([]byte, 256*1024)
	lastReport := int64(0)
	for {
		nr, er := resp.Body.Read(buf)
		if nr > 0 {
			wn, ew := f.Write(buf[:nr])
			n += int64(wn)
			if ew != nil {
				err = ew
				break
			}
			if n-lastReport >= 1024*1024 {
				progress := 0.0
				if resp.ContentLength > 0 {
					progress = float64(n) * 100 / float64(resp.ContentLength)
				}
				a.db.Model(&t).Updates(map[string]any{"downloaded_bytes": n, "total_bytes": resp.ContentLength, "progress": progress})
				lastReport = n
			}
		}
		if er == io.EOF {
			break
		}
		if er != nil {
			err = er
			break
		}
	}
	cerr := f.Close()
	if err != nil {
		a.fail(&t, err)
		return
	}
	if cerr != nil {
		a.fail(&t, cerr)
		return
	}
	if t.TotalBytes > 0 && n != t.TotalBytes {
		a.fail(&t, fmt.Errorf("incomplete download: got %d of %d bytes", n, t.TotalBytes))
		return
	}
	actual, detectErr := detectAudioFormat(tmp)
	if detectErr != nil {
		a.fail(&t, fmt.Errorf("detect audio format: %w", detectErr))
		_ = os.Remove(tmp)
		return
	}
	actual = normalizedFormat(actual)
	local = strings.TrimSuffix(local, filepath.Ext(local)) + "." + actual
	if err = os.Rename(tmp, local); err != nil {
		a.fail(&t, err)
		_ = os.Remove(tmp)
		return
	}
	t.Format = actual
	t.LocalPath = local
	t.ActualQuality = actualQuality(t.Quality, t.ResolvedQuality, actual)
	t.DownloadedBytes = n
	if t.TotalBytes > 0 {
		t.Progress = 100
	}
	a.db.Save(&t)
	a.upload(&t)
}

func (a *App) resolve(t Task) (playInfo, error) {
	u, _ := url.Parse(strings.TrimRight(a.cfg.AMAPI, "/") + "/song/play-url")
	q := u.Query()
	q.Set("source", t.Source)
	q.Set("rid", t.RID)
	q.Set("level", t.Quality)
	u.RawQuery = q.Encode()
	resp, err := a.httpClient.Get(u.String())
	if err != nil {
		return playInfo{}, err
	}
	defer resp.Body.Close()
	var x struct {
		Code    int      `json:"code"`
		Data    playInfo `json:"data"`
		Message string   `json:"message"`
	}
	if err = json.NewDecoder(resp.Body).Decode(&x); err != nil {
		return playInfo{}, err
	}
	if x.Code != 0 || x.Data.URL == "" {
		return playInfo{}, errors.New(x.Message)
	}
	return x.Data, nil
}

func (a *App) upload(t *Task) {
	if a.cfg.WebDAVURL == "" {
		t.Status = "done"
		a.db.Save(t)
		return
	}
	a.db.Model(t).Updates(map[string]any{"status": "uploading"})
	remote := strings.TrimRight(a.cfg.WebDAVPath, "/") + "/" + safe(t.Artist) + "/" + safe(t.Album) + "/" + safe(t.Title) + "." + safe(t.Format)
	t.RemotePath = remote
	remoteDir := filepath.ToSlash(filepath.Dir(remote))
	client := gowebdav.NewClient(a.cfg.WebDAVURL, a.cfg.WebDAVUser, a.cfg.WebDAVPass)
	client.SetTimeout(10 * time.Minute)
	if err := client.MkdirAll(remoteDir, 0755); err != nil {
		a.fail(t, err)
		return
	}
	if fi, err := os.Stat(t.LocalPath); err == nil {
		if ri, err := client.Stat(remote); err == nil && ri.Size() == fi.Size() {
			t.Status = "done"
			t.Progress = 100
			a.db.Save(t)
			return
		}
	}
	f, err := os.Open(t.LocalPath)
	if err != nil {
		a.fail(t, err)
		return
	}
	defer f.Close()
	if err = client.WriteStream(remote, f, 0644); err != nil {
		a.fail(t, err)
		return
	}
	t.Status = "done"
	t.Progress = 100
	a.db.Save(t)
}
func (a *App) fail(t *Task, err error) {
	if t.LocalPath != "" {
		_ = os.Remove(t.LocalPath + ".part")
	}
	t.Status = "error"
	t.Error = err.Error()
	a.db.Save(t)
}

func (a *App) getSettings(c *gin.Context) {
	c.JSON(200, gin.H{"amApi": a.cfg.AMAPI, "downloadDir": a.cfg.DownloadDir, "webdavUrl": a.cfg.WebDAVURL, "webdavUser": a.cfg.WebDAVUser, "webdavPass": a.cfg.WebDAVPass, "webdavPath": a.cfg.WebDAVPath, "quality": a.cfg.DefaultQuality})
}
func (a *App) putSettings(c *gin.Context) {
	var v map[string]string
	if c.BindJSON(&v) != nil {
		c.Status(400)
		return
	}
	if x := v["amApi"]; x != "" {
		a.cfg.AMAPI = x
		a.save("am_api", x)
	}
	if x := v["downloadDir"]; x != "" {
		a.cfg.DownloadDir = x
		a.save("download_dir", x)
		os.MkdirAll(x, 0755)
	}
	a.cfg.WebDAVURL = v["webdavUrl"]
	a.cfg.WebDAVUser = v["webdavUser"]
	a.cfg.WebDAVPass = v["webdavPass"]
	a.cfg.WebDAVPath = strings.Trim(v["webdavPath"], "/")
	a.cfg.DefaultQuality = v["quality"]
	for k, x := range map[string]string{"webdav_url": a.cfg.WebDAVURL, "webdav_user": a.cfg.WebDAVUser, "webdav_pass": a.cfg.WebDAVPass, "webdav_path": a.cfg.WebDAVPath, "quality": a.cfg.DefaultQuality} {
		a.save(k, x)
	}
	c.JSON(200, gin.H{"ok": true})
}
func (a *App) testWebDAV(c *gin.Context) {
	var v struct{ URL, User, Pass, Path string }
	if c.BindJSON(&v) != nil {
		c.Status(400)
		return
	}
	client := gowebdav.NewClient(v.URL, v.User, v.Pass)
	client.SetTimeout(20 * time.Second)
	if err := client.Connect(); err != nil {
		c.JSON(502, gin.H{"ok": false, "error": err.Error()})
		return
	}
	c.JSON(200, gin.H{"ok": true, "path": v.Path})
}

var bad = regexp.MustCompile(`[/\\:*?"<>|]+`)

func safe(s string) string {
	s = strings.TrimSpace(s)
	s = bad.ReplaceAllString(s, "_")
	s = strings.Trim(s, ". ")
	if s == "" {
		return "Unknown"
	}
	return s
}
func getenv(k, d string) string {
	if v := os.Getenv(k); v != "" {
		return v
	}
	return d
}
func atoi(s string, d int) int {
	v, e := strconv.Atoi(s)
	if e != nil || v < 1 {
		return d
	}
	return v
}
func must(e error) {
	if e != nil {
		panic(e)
	}
}