package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"
)

const (
	tidalBaseURL  = "https://api.tidal.com/v1"
	tidalAuthURL  = "https://auth.tidal.com/v1/oauth2/token"
	defaultPort   = ":8080"
)

type Config struct {
	Port            string
	RequiredUA      string
	AllowAnyQuality bool
	CountryCode     string
	TidalClientID   string
	TidalSecret     string
}

type TrackRequest struct {
	ID      string `json:"id"`
	Quality string `json:"quality"`
}

type TidalPlaybackInfo struct {
	TrackID              int           `json:"trackId"`
	AssetPresentation    string        `json:"assetPresentation"`
	AudioQuality         string        `json:"audioQuality"`
	AudioMode            string        `json:"audioMode"`
	ManifestMimeType     string        `json:"manifestMimeType"`
	ManifestHash         string        `json:"manifestHash"`
	Manifest             string        `json:"manifest"`
	BitDepth             int           `json:"bitDepth,omitempty"`
	SampleRate           int           `json:"sampleRate,omitempty"`
	ReplayGain           float64       `json:"replayGain"`
	AlbumReplayGain      float64       `json:"albumReplayGain"`
	TrackPeakAmplitude   float64       `json:"trackPeakAmplitude"`
	AlbumPeakAmplitude   float64       `json:"albumPeakAmplitude"`
	DRMData              interface{}   `json:"drmData"`
	Formats              []interface{} `json:"formats"`
}

type ProxyResponse struct {
	Success bool   `json:"success"`
	Data    string `json:"data,omitempty"`
	Error   string `json:"error,omitempty"`
}

type tokenCache struct {
	mu        sync.RWMutex
	token     string
	expiresAt time.Time
}

var (
	config       Config
	client       *http.Client
	startTime    time.Time
	appTokenCache tokenCache
)

func loadConfig() Config {
	return Config{
		Port:            getEnv("PORT", defaultPort),
		RequiredUA:      getEnv("REQUIRED_UA_PREFIX", "SpotiFLAC-Mobile/"),
		AllowAnyQuality: getEnv("ALLOW_ANY_QUALITY", "false") == "true",
		CountryCode:     getEnv("COUNTRY_CODE", "US"),
		TidalClientID:   getEnv("TIDAL_CLIENT_ID", ""),
		TidalSecret:     getEnv("TIDAL_CLIENT_SECRET", ""),
	}
}

func getEnv(key, fallback string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return fallback
}

func validateUA(r *http.Request) bool {
	ua := r.Header.Get("User-Agent")
	if ua == "" {
		return false
	}
	return strings.HasPrefix(ua, config.RequiredUA)
}

func getBearerToken(r *http.Request) string {
	if auth := r.Header.Get("Authorization"); auth != "" {
		return strings.TrimPrefix(auth, "Bearer ")
	}
	if token := r.Header.Get("X-Tidal-Token"); token != "" {
		return token
	}
	return r.URL.Query().Get("token")
}

func getAppToken() (string, error) {
	appTokenCache.mu.RLock()
	if appTokenCache.token != "" && time.Now().Before(appTokenCache.expiresAt) {
		token := appTokenCache.token
		appTokenCache.mu.RUnlock()
		return token, nil
	}
	appTokenCache.mu.RUnlock()

	appTokenCache.mu.Lock()
	defer appTokenCache.mu.Unlock()

	if appTokenCache.token != "" && time.Now().Before(appTokenCache.expiresAt) {
		return appTokenCache.token, nil
	}

	if config.TidalClientID != "" && config.TidalSecret != "" {
		token, expiresAt, err := fetchAppToken(config.TidalClientID, config.TidalSecret)
		if err == nil {
			appTokenCache.token = token
			appTokenCache.expiresAt = expiresAt
			return token, nil
		}
		log.Printf("Failed to get app token with configured credentials: %v", err)
	}

	// Fallback: try known Tidal client credentials
	clients := []struct {
		id     string
		secret string
	}{
		// PKCE credentials (Hi-Res capable) from tidalapi — may still be valid
		{"6BDSRdpK9hqEBTgU", "xeuPmY7nbpZ9IIbLAcQ93shka1VNheUAqN6IcszjTG8="},
		{"fX2JxdmntZWK0ixT", "1Nn9AfDAjxrgJFJbKNWLeAyKGVGmINuXPPLHVXAvxAg="},
		{"7m7Ap0JC9j1cOM3n", "caDRxBDSx9WtZz7nH9Mz8uSiVn6FL7dP"},
		{"zU4XHVVkc2tDPo4t", "VJKhYqM4Cj5z1TNnX1hB8eKhE3w7LxvO9Qp2RsJ6"},
	}

	for _, c := range clients {
		token, expiresAt, err := fetchAppToken(c.id, c.secret)
		if err == nil {
			appTokenCache.token = token
			appTokenCache.expiresAt = expiresAt
			log.Printf("Got app token with client %s", c.id[:4]+"...")
			return token, nil
		}
	}

	return "", fmt.Errorf("all client credentials failed")
}

func fetchAppToken(clientID, secret string) (string, time.Time, error) {
	creds := base64.StdEncoding.EncodeToString([]byte(clientID + ":" + secret))

	req, err := http.NewRequest("POST", tidalAuthURL, strings.NewReader("grant_type=client_credentials"))
	if err != nil {
		return "", time.Time{}, err
	}
	req.Header.Set("Authorization", "Basic "+creds)
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := client.Do(req)
	if err != nil {
		return "", time.Time{}, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return "", time.Time{}, fmt.Errorf("auth failed: %d %s", resp.StatusCode, string(body))
	}

	var result struct {
		AccessToken string `json:"access_token"`
		ExpiresIn   int    `json:"expires_in"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return "", time.Time{}, err
	}

	return result.AccessToken, time.Now().Add(time.Duration(result.ExpiresIn-60) * time.Second), nil
}

func callTidalAPI(endpoint string, token string, params map[string]string) (*http.Response, error) {
	req, err := http.NewRequest("GET", endpoint, nil)
	if err != nil {
		return nil, err
	}

	q := req.URL.Query()
	q.Set("countryCode", config.CountryCode)
	for k, v := range params {
		q.Set(k, v)
	}
	req.URL.RawQuery = q.Encode()

	if token != "" {
		req.Header.Set("Authorization", "Bearer "+token)
	} else {
		appToken, err := getAppToken()
		if err != nil {
			return nil, fmt.Errorf("no valid token: %w", err)
		}
		req.Header.Set("Authorization", "Bearer "+appToken)
	}

	req.Header.Set("User-Agent", "TIDAL_ANDROID/1479 okhttp/3.14.9")

	return client.Do(req)
}

func corsMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type, User-Agent, X-Tidal-Token, Authorization")

		if r.Method == "OPTIONS" {
			w.WriteHeader(http.StatusOK)
			return
		}

		next(w, r)
	}
}

func handleDownloadTrack(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeJSON(w, http.StatusMethodNotAllowed, ProxyResponse{Success: false, Error: "Method not allowed"})
		return
	}

	if !validateUA(r) {
		writeJSON(w, http.StatusForbidden, ProxyResponse{
			Success: false,
			Error:   "Invalid user agent. Use the original app User-Agent in the format AppName/Version",
		})
		return
	}

	var req TrackRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeJSON(w, http.StatusBadRequest, ProxyResponse{Success: false, Error: "Invalid request body"})
		return
	}

	if req.ID == "" {
		writeJSON(w, http.StatusBadRequest, ProxyResponse{Success: false, Error: "Track ID is required"})
		return
	}

	userToken := getBearerToken(r)
	quality := normalizeQuality(req.Quality)

	// Try direct Tidal API first (uses fallback credentials if needed)
	playbackInfo, apiErr := getPlaybackInfo(req.ID, quality, userToken)
	if apiErr == nil {
		// Reject PREVIEW (30s demo) — fall through to zarz.moe for full track
		if strings.EqualFold(playbackInfo.AssetPresentation, "PREVIEW") {
			log.Printf("Direct Tidal API returned PREVIEW for track %s, falling back to api.zarz.moe", req.ID)
		} else {
			writeResponse(w, playbackInfo)
			return
		}
	} else {
		log.Printf("Direct Tidal API failed for track %s: %v, falling back to api.zarz.moe", req.ID, apiErr)
	}

	// Fallback: forward to api.zarz.moe with retries
	var lastErr error
	for attempt := 0; attempt < 3; attempt++ {
		if attempt > 0 {
			time.Sleep(time.Duration(attempt) * time.Second)
		}
		respBody, statusCode, err := forwardToZarz(req.ID, quality)
		if err == nil {
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(statusCode)
			w.Write(respBody)
			return
		}
		lastErr = err
		// Don't retry on rate limits (retry_after is typically long)
		if strings.Contains(err.Error(), "cooldown") || strings.Contains(err.Error(), "retry_after") {
			break
		}
	}

	log.Printf("All zarz.moe attempts failed for track %s: %v", req.ID, lastErr)
	writeJSON(w, http.StatusBadGateway, ProxyResponse{
		Success: false,
		Error:   fmt.Sprintf("Failed to forward: %v", lastErr),
	})
}

func forwardToZarz(trackID, quality string) ([]byte, int, error) {
	body := map[string]interface{}{
		"id":      trackID,
		"quality": quality,
	}
	bodyBytes, _ := json.Marshal(body)

	req, err := http.NewRequest("POST", "https://api.zarz.moe/v1/dl/tid2", bytes.NewReader(bodyBytes))
	if err != nil {
		return nil, 0, err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("User-Agent", "SpotiFLAC-Mobile/4.5.5")

	resp, err := client.Do(req)
	if err != nil {
		return nil, 0, err
	}
	defer resp.Body.Close()

	respBody, _ := io.ReadAll(resp.Body)

	if resp.StatusCode != http.StatusOK {
		// Try to parse as the new format (success field)
		var result struct {
			Success bool   `json:"success"`
			Error   string `json:"error"`
		}
		if json.Unmarshal(respBody, &result) == nil && !result.Success && result.Error != "" {
			return nil, resp.StatusCode, fmt.Errorf("api.zarz.moe: %s", result.Error)
		}
		return nil, resp.StatusCode, fmt.Errorf("api.zarz.moe returned %d: %s", resp.StatusCode, string(respBody))
	}

	return respBody, resp.StatusCode, nil
}

func writeResponse(w http.ResponseWriter, info *TidalPlaybackInfo) {
	if !config.AllowAnyQuality {
		rq := info.AudioQuality
		if rq == "LOW" || rq == "HIGH" {
			writeJSON(w, http.StatusBadGateway, ProxyResponse{
				Success: false,
				Error:   fmt.Sprintf("Requested lossless but got %s. Track may not support it.", rq),
			})
			return
		}
	}

	dataBytes, err := json.Marshal(info)
	if err != nil {
		writeJSON(w, http.StatusInternalServerError, ProxyResponse{Success: false, Error: "Failed to encode response"})
		return
	}

	writeJSON(w, http.StatusOK, ProxyResponse{Success: true, Data: string(dataBytes)})
}

func normalizeQuality(quality string) string {
	switch strings.ToUpper(quality) {
	case "HI_RES_LOSSLESS", "HIRES_LOSSLESS", "MASTER", "HI_RES", "HIRES":
		return "HI_RES_LOSSLESS"
	case "LOSSLESS", "FLAC", "HIFI":
		return "LOSSLESS"
	case "HIGH", "NORMAL":
		return "HIGH"
	case "LOW":
		return "LOW"
	default:
		return "LOSSLESS"
	}
}

func getPlaybackInfo(trackID, quality, token string) (*TidalPlaybackInfo, error) {
	params := map[string]string{
		"playbackMode":      "STREAM",
		"audioquality":      quality,
		"assetpresentation": "FULL",
	}

	resp, err := callTidalAPI(
		fmt.Sprintf("%s/tracks/%s/playbackinfo", tidalBaseURL, trackID),
		token,
		params,
	)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("Tidal API returned %d: %s", resp.StatusCode, string(body))
	}

	var info TidalPlaybackInfo
	if err := json.NewDecoder(resp.Body).Decode(&info); err != nil {
		return nil, err
	}

	if info.Formats == nil {
		info.Formats = extractFormatsFromManifest(info.Manifest, info.ManifestMimeType)
	}

	return &info, nil
}

func extractFormatsFromManifest(manifest, mimeType string) []interface{} {
	formats := make([]interface{}, 0)
	if mimeType == "" || manifest == "" {
		return formats
	}

	decoded, err := base64.StdEncoding.DecodeString(manifest)
	if err != nil {
		return formats
	}

	content := string(decoded)
	if !strings.Contains(content, "<MPD") {
		return formats
	}

	if strings.Contains(content, "codecs=\"fLaC\"") || strings.Contains(content, "codecs=\"flac\"") {
		if strings.Contains(content, "HI_RES") || strings.Contains(content, "MAX_SAMPLE_RATE") {
			formats = append(formats, "FLAC_HIRES")
		}
		formats = append(formats, "FLAC")
	}
	if strings.Contains(content, "codecs=\"ec-3\"") {
		formats = append(formats, "EAC3_JOC")
	}
	if strings.Contains(content, "codecs=\"mp4a.40.2\"") {
		formats = append(formats, "AACLC")
	}
	if strings.Contains(content, "codecs=\"mp4a.40.5\"") {
		formats = append(formats, "HEAACV1")
	}

	return formats
}

func handleHealth(w http.ResponseWriter, r *http.Request) {
	hasToken := config.TidalClientID != "" || appTokenCache.token != ""
	writeJSON(w, http.StatusOK, map[string]interface{}{
		"status":   "ok",
		"uptime":   time.Since(startTime).String(),
		"hasToken": hasToken,
	})
}

func handleAuth(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeJSON(w, http.StatusMethodNotAllowed, ProxyResponse{Success: false, Error: "Send POST with Tidal credentials"})
		return
	}

	var creds struct {
		ClientID     string `json:"client_id"`
		ClientSecret string `json:"client_secret"`
	}
	if err := json.NewDecoder(r.Body).Decode(&creds); err != nil {
		writeJSON(w, http.StatusBadRequest, ProxyResponse{Success: false, Error: "Invalid request"})
		return
	}

	if creds.ClientID != "" && creds.ClientSecret != "" {
		token, expiresAt, err := fetchAppToken(creds.ClientID, creds.ClientSecret)
		if err != nil {
			writeJSON(w, http.StatusBadGateway, ProxyResponse{Success: false, Error: err.Error()})
			return
		}

		appTokenCache.mu.Lock()
		appTokenCache.token = token
		appTokenCache.expiresAt = expiresAt
		appTokenCache.mu.Unlock()

		config.TidalClientID = creds.ClientID
		config.TidalSecret = creds.ClientSecret

		writeJSON(w, http.StatusOK, ProxyResponse{Success: true, Data: "Credentials accepted"})
	} else {
		writeJSON(w, http.StatusBadRequest, ProxyResponse{Success: false, Error: "client_id and client_secret required"})
	}
}

type QobuzRequest struct {
	TrackID string `json:"trackId"`
	Quality string `json:"quality"`
}

func handleDownloadQobuz(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeJSON(w, http.StatusMethodNotAllowed, ProxyResponse{Success: false, Error: "Method not allowed"})
		return
	}

	if !validateUA(r) {
		writeJSON(w, http.StatusForbidden, ProxyResponse{
			Success: false,
			Error:   "Invalid user agent",
		})
		return
	}

	trackID := r.URL.Query().Get("trackId")
	if trackID == "" {
		writeJSON(w, http.StatusBadRequest, ProxyResponse{Success: false, Error: "trackId is required"})
		return
	}

	quality := r.URL.Query().Get("quality")
	if quality == "" {
		quality = "HI_RES_LOSSLESS"
	}

	req, err := http.NewRequest("GET", "https://api.zarz.moe/v1/dl/qbz", r.Body)
	if err != nil {
		writeJSON(w, http.StatusBadGateway, ProxyResponse{Success: false, Error: err.Error()})
		return
	}

	q := req.URL.Query()
	q.Set("trackId", trackID)
	q.Set("quality", quality)
	req.URL.RawQuery = q.Encode()

	req.Header.Set("User-Agent", r.Header.Get("User-Agent"))
	req.Header.Set("Accept", "application/json")

	resp, err := client.Do(req)
	if err != nil {
		writeJSON(w, http.StatusBadGateway, ProxyResponse{Success: false, Error: err.Error()})
		return
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(resp.Body)
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(resp.StatusCode)
	w.Write(body)
}

type AmazonRequest struct {
	ASIN  string `json:"asin"`
	Codec string `json:"codec"`
}

func handleDownloadAmazon(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeJSON(w, http.StatusMethodNotAllowed, ProxyResponse{Success: false, Error: "Method not allowed"})
		return
	}

	if !validateUA(r) {
		writeJSON(w, http.StatusForbidden, ProxyResponse{
			Success: false,
			Error:   "Invalid user agent",
		})
		return
	}

	asin := r.URL.Query().Get("asin")
	if asin == "" {
		writeJSON(w, http.StatusBadRequest, ProxyResponse{Success: false, Error: "asin is required"})
		return
	}

	codec := r.URL.Query().Get("codec")
	if codec == "" {
		codec = "flac"
	}

	req, err := http.NewRequest("GET", "https://api.zarz.moe/v1/dl/amazeamazeamaze/media", nil)
	if err != nil {
		writeJSON(w, http.StatusBadGateway, ProxyResponse{Success: false, Error: err.Error()})
		return
	}

	q := req.URL.Query()
	q.Set("asin", asin)
	q.Set("codec", codec)
	req.URL.RawQuery = q.Encode()

	req.Header.Set("User-Agent", r.Header.Get("User-Agent"))
	req.Header.Set("Accept", "application/json")

	resp, err := client.Do(req)
	if err != nil {
		writeJSON(w, http.StatusBadGateway, ProxyResponse{Success: false, Error: err.Error()})
		return
	}
	defer resp.Body.Close()

	body, _ := io.ReadAll(resp.Body)
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(resp.StatusCode)
	w.Write(body)
}

func handleProxy(w http.ResponseWriter, r *http.Request) {
	targetURL := r.URL.Query().Get("url")
	if targetURL == "" {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "url parameter required"})
		return
	}

	req, err := http.NewRequest(r.Method, targetURL, r.Body)
	if err != nil {
		writeJSON(w, http.StatusBadGateway, map[string]string{"error": err.Error()})
		return
	}

	// Forward relevant headers
	for _, h := range []string{"Range", "If-None-Match", "If-Modified-Since"} {
		if v := r.Header.Get(h); v != "" {
			req.Header.Set(h, v)
		}
	}

	resp, err := client.Do(req)
	if err != nil {
		writeJSON(w, http.StatusBadGateway, map[string]string{"error": err.Error()})
		return
	}
	defer resp.Body.Close()

	// Copy response headers
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.Header().Set("Access-Control-Expose-Headers", "Content-Length, Content-Range, Accept-Ranges")
	for _, h := range []string{"Content-Type", "Content-Length", "Content-Range", "Accept-Ranges", "Cache-Control", "ETag"} {
		if v := resp.Header.Get(h); v != "" {
			w.Header().Set(h, v)
		}
	}

	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}

func writeJSON(w http.ResponseWriter, status int, v interface{}) {
	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(v)
}

func main() {
	config = loadConfig()
	startTime = time.Now()

	client = &http.Client{
		Timeout: 60 * time.Second,
	}

	// Try to get an initial app token on startup
	go func() {
		if config.TidalClientID != "" && config.TidalSecret != "" {
			token, expiresAt, err := fetchAppToken(config.TidalClientID, config.TidalSecret)
			if err == nil {
				appTokenCache.mu.Lock()
				appTokenCache.token = token
				appTokenCache.expiresAt = expiresAt
				appTokenCache.mu.Unlock()
				log.Println("Initial app token acquired")
			} else {
				log.Printf("Failed to get initial app token: %v", err)
			}
		}
	}()

	mux := http.NewServeMux()
	mux.HandleFunc("/v1/dl/tid2", corsMiddleware(handleDownloadTrack))
	mux.HandleFunc("/v1/dl/qbz", corsMiddleware(handleDownloadQobuz))
	mux.HandleFunc("/v1/dl/amazeamazeamaze/media", corsMiddleware(handleDownloadAmazon))
	mux.HandleFunc("/health", corsMiddleware(handleHealth))
	mux.HandleFunc("/auth", corsMiddleware(handleAuth))
	mux.HandleFunc("/proxy", corsMiddleware(handleProxy))

	log.Printf("Starting Tidal proxy server on %s", config.Port)
	log.Printf("Required UA prefix: %s", config.RequiredUA)
	log.Printf("Allow any quality: %v", config.AllowAnyQuality)
	log.Printf("Country code: %s", config.CountryCode)
	log.Printf("Has Tidal credentials: %v", config.TidalClientID != "")

	if err := http.ListenAndServe(config.Port, mux); err != nil {
		log.Fatalf("Server failed: %v", err)
	}
}