package main

import (
	"encoding/json"
	"encoding/xml"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return err == nil
}

type MPD struct {
	XMLName xml.Name `xml:"MPD"`
	PSSHs   []string `xml:"Period>AdaptationSet>ContentProtection>PSSH"`
}

func GetDRMMPDPSSH(mpdURL, policy, signature, kvp string) (string, error) {
	req, _ := http.NewRequest("GET", mpdURL, nil)
	req.Header.Set("Accept", "application/dash+xml,video/vnd.mpeg.dash.mpd")
	req.Header.Set("User-Agent", "Mozilla/5.0")
	req.Header.Set("Cookie", fmt.Sprintf("CloudFront-Policy=%s; CloudFront-Signature=%s; CloudFront-Key-Pair-Id=%s", policy, signature, kvp))

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

	xmlData, _ := io.ReadAll(resp.Body)
	var mpd MPD
	if err := xml.Unmarshal(xmlData, &mpd); err != nil {
		return "", err
	}
	if len(mpd.PSSHs) == 0 {
		return "", errors.New("PSSH not found")
	}
	return mpd.PSSHs[0], nil
}

func GetDRMMPDLastModified(mpdURL, policy, signature, kvp string) (string, error) {
	req, _ := http.NewRequest("HEAD", mpdURL, nil)
	req.Header.Set("Accept", "application/dash+xml,video/vnd.mpeg.dash.mpd")
	req.Header.Set("User-Agent", "Mozilla/5.0")
	req.Header.Set("Cookie", fmt.Sprintf("CloudFront-Policy=%s; CloudFront-Signature=%s; CloudFront-Key-Pair-Id=%s", policy, signature, kvp))

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

	lastModified := resp.Header.Get("Last-Modified")
	return lastModified, nil
}

func GetDecryptionKeyOFDL(drmHeaders map[string]string, licenseURL, pssh string) (string, error) {
	form := url.Values{}
	headersJSON, _ := json.Marshal(drmHeaders)

	form.Set("license_url", licenseURL)
	form.Set("headers", string(headersJSON))
	form.Set("pssh", pssh)
	form.Set("build_identifier", "windows_software_widevinecdm_win_x86_64")
	form.Set("proxy", "")

	resp, err := http.PostForm("", form)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	var result struct {
		Keys []struct {
			Key string `json:"key"`
		} `json:"keys"`
	}

	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return "", err
	}
	if len(result.Keys) == 0 {
		return "", errors.New("no keys returned")
	}
	return result.Keys[0].Key, nil
}

// ⬇ Replace with your own CDM client logic
func GetDecryptionKeyCDM(drmHeaders map[string]string, licenseURL, pssh string) (string, error) {
	// This is a stub. Replace with your CDM library logic.
	// Example:
	// return widevine.GetContentKey(licenseURL, drmHeaders, pssh)
	return "", errors.New("local CDM support not implemented")
}

func DownloadPurchasedMessageDRMVideo(mpdURL, decryptionKey, outputDir, filename string) error {
	outPath := filepath.Join(outputDir, filename)
	if fileExists(outPath) {
		fmt.Println("Already downloaded:", outPath)
		return nil
	}

	cmd := exec.Command("shaka-packager",
		fmt.Sprintf("input=%s,stream=video,output=%s", mpdURL, outPath),
		fmt.Sprintf("--keys=key_id=00000000000000000000000000000000:key=%s", decryptionKey),
		"--allow-multiple-key-ids",
		"--quiet",
	)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func DownloadPurchasedMedia(url, outputDir, filename string) error {
	outPath := filepath.Join(outputDir, filename)
	if fileExists(outPath) {
		fmt.Println("Already downloaded:", outPath)
		return nil
	}

	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	out, err := os.Create(outPath)
	if err != nil {
		return err
	}
	defer out.Close()

	fmt.Println("Downloading:", outPath)
	_, err = io.Copy(out, resp.Body)
	return err
}

// 🧠 Smart filename formatting
func GenerateFilename(mediaID, userID, msgID string, isDRM bool) string {
	if isDRM {
		return fmt.Sprintf("drm_%s_%s_%s.mp4", userID, msgID, mediaID)
	}
	return fmt.Sprintf("media_%s_%s_%s.mp4", userID, msgID, mediaID)
}

// 🧠 Dummy auth header generator
func GetDynamicHeaders(endpoint, query string) map[string]string {
	return map[string]string{
		"User-Agent": "Mozilla/5.0",
		"Cookie":     "sess=your_cookie_here",
	}
}

func main() {
	// Example data — Replace with real input
	paidMessageValue := "https://cdn3.onlyfans.com/dash/files,..."
	paidMessageKey := "mediaId123"
	outputDir := "./downloads"
	deviceFolder := "./device"
	deviceName := "default"
	hasSelectedUsers := true

	clientIdBlobMissing := !fileExists(fmt.Sprintf("%s/%s/device_client_id_blob", deviceFolder, deviceName))
	devicePrivateKeyMissing := !fileExists(fmt.Sprintf("%s/%s/device_private_key", deviceFolder, deviceName))

	err := os.MkdirAll(outputDir, 0755)
	if err != nil {
		panic(err)
	}

	if strings.Contains(paidMessageValue, "cdn3.onlyfans.com/dash/files") {
		parts := strings.Split(paidMessageValue, ",")
		if len(parts) < 6 {
			fmt.Println("Invalid media info string.")
			return
		}

		mpdURL := parts[0]
		policy := parts[1]
		signature := parts[2]
		kvp := parts[3]
		mediaID := parts[4]
		messageID := parts[5]
		userID := "userXYZ" // Replace with actual username

		pssh, err := GetDRMMPDPSSH(mpdURL, policy, signature, kvp)
		if err != nil {
			fmt.Println("Failed to get PSSH:", err)
			return
		}

		drmHeaders := GetDynamicHeaders(fmt.Sprintf("/api2/v2/users/media/%s/drm/message/%s", mediaID, messageID), "?type=widevine")

		var decryptionKey string
		if clientIdBlobMissing || devicePrivateKeyMissing {
			decryptionKey, err = GetDecryptionKeyOFDL(drmHeaders, fmt.Sprintf("https://onlyfans.com/api2/v2/users/media/%s/drm/message/%s?type=widevine", mediaID, messageID), pssh)
		} else {
			decryptionKey, err = GetDecryptionKeyCDM(drmHeaders, fmt.Sprintf("https://onlyfans.com/api2/v2/users/media/%s/drm/message/%s?type=widevine", mediaID, messageID), pssh)
		}

		if err != nil {
			fmt.Println("Failed to get decryption key:", err)
			return
		}

		filename := GenerateFilename(mediaID, userID, messageID, true)
		err = DownloadPurchasedMessageDRMVideo(mpdURL, decryptionKey, outputDir, filename)
		if err != nil {
			fmt.Println("Failed to download DRM video:", err)
		}
	} else {
		userID := "userXYZ"
		msgID := "msg456"
		filename := GenerateFilename(paidMessageKey, userID, msgID, false)
		err := DownloadPurchasedMedia(paidMessageValue, outputDir, filename)
		if err != nil {
			fmt.Println("Failed to download non-DRM video:", err)
		}
	}
}