package main

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

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

type Config struct {
	Port            string
	RequiredUA      string
	AllowAnyQuality bool
	CountryCode     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"`
}

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

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"),
	}
}

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 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()

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

	return client.Do(req)
}

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
	}

	token := getBearerToken(r)
	if token == "" {
		writeJSON(w, http.StatusUnauthorized, ProxyResponse{
			Success: false,
			Error:   "No Tidal token provided. Send your Tidal access token via Authorization: Bearer <token> or X-Tidal-Token header",
		})
		return
	}

	quality := normalizeQuality(req.Quality)

	playbackInfo, err := getPlaybackInfo(req.ID, quality, token)
	if err != nil {
		log.Printf("Error getting playback info for track %s: %v", req.ID, err)
		writeJSON(w, http.StatusBadGateway, ProxyResponse{
			Success: false,
			Error:   fmt.Sprintf("Failed to get playback info: %v", err),
		})
		return
	}

	if !config.AllowAnyQuality {
		rq := playbackInfo.AudioQuality
		if rq == "LOW" || rq == "HIGH" {
			writeJSON(w, http.StatusBadGateway, ProxyResponse{
				Success: false,
				Error:   fmt.Sprintf("Requested %s quality but Tidal returned %s. This may indicate the track doesn't support lossless or your account lacks premium access.", quality, rq),
			})
			return
		}
	}

	dataBytes, err := json.Marshal(playbackInfo)
	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,
	}

	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) {
	writeJSON(w, http.StatusOK, map[string]string{
		"status": "ok",
		"uptime": time.Since(startTime).String(),
	})
}

func handleAuth(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, http.StatusOK, map[string]string{
		"status":  "ok",
		"message": "Send your Tidal token via Authorization: Bearer <token> or X-Tidal-Token header",
	})
}

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

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

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

	mux := http.NewServeMux()
	mux.HandleFunc("/v1/dl/tid2", handleDownloadTrack)
	mux.HandleFunc("/health", handleHealth)
	mux.HandleFunc("/auth", handleAuth)

	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)

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