// Package phonecap runs one phone MITM capture: proxy → launch → autoplay →
// wait → wvkey → session → optional force-stop. Used by capture.exe and agent.
package phonecap

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"drmdecryption/adb"
	"drmdecryption/app"
	"drmdecryption/capture"
	"drmdecryption/proxy"
	"drmdecryption/repo"
	"drmdecryption/session"
	"drmdecryption/wvkey"
)

// DefaultCAHash is the subject hash of the bundled MITM CA, used when an app
// module does not state its own.
const DefaultCAHash = "6c3578b4"

// Options for a single capture job on one device.
type Options struct {
	App        app.App
	Channel    string
	Client     *adb.Client // serial-scoped when multi-device
	Root       string
	Python     string
	WVD        string // path to .wvd device file (required for key fetch)
	UserAgent  string // optional license-request User-Agent
	Wait       time.Duration
	NoAutoPlay bool
	// KeyMode: "auto" (default), "modulardrm", or "raw".
	KeyMode string
	// CloseApp force-stops the package when the job ends (success or failure).
	CloseApp bool
	// ClearProxy clears global http_proxy after the job.
	ClearProxy bool
}

// Result is the captured session plus on-disk path.
type Result struct {
	Session session.Session
	Path    string
	Key     string
	MPD     string
}

// RunPassive starts the MITM only — no app launch / auto-play. Waits for
// whatever DRM/manifest traffic the phone generates, prints it, and writes
// outputs/capture/<stamp>/ when anything useful appears.
func RunPassive(opt Options) (res Result, err error) {
	if opt.Client == nil {
		opt.Client = adb.New()
	}
	if opt.Root == "" {
		opt.Root = repo.Root()
	}
	if opt.Wait <= 0 {
		opt.Wait = 180 * time.Second
	}
	c := opt.Client
	root := opt.Root

	if opt.ClearProxy {
		defer ClearProxyLater(c)
	}
	if err := c.EnsureDevice(); err != nil {
		return res, err
	}

	proxyBin, caHash := resolveProxy(nil, root)
	if err := proxy.PushAndStart(c, proxyBin); err != nil {
		return res, err
	}
	if err := proxy.EnsureHTTPProxy(c, ""); err != nil {
		return res, err
	}
	proxy.ReinjectCA(c, caHash)

	fmt.Println("[*] Passive capture — open any app and start playback on the phone")
	fmt.Println("    (no auto-play; Ctrl+C will not save — wait for traffic or raise --wait)")

	hints := capture.DefaultPassiveHints()
	_ = os.MkdirAll(filepath.Join(root, ".cache"), 0o755)
	serialTag := c.Serial
	if serialTag == "" {
		serialTag = "default"
	}
	localLog := filepath.Join(root, ".cache", "appproxy-"+sanitize(serialTag)+".log")
	localCap := filepath.Join(root, ".cache", "appproxy-"+sanitize(serialTag)+"_cap.json")
	stop := make(chan struct{})
	go capture.MirrorLogLoop(c, hints.RemoteLog, localLog, stop)
	defer close(stop)

	fmt.Println("[*] Waiting for any license / PSSH / manifest…")
	fmt.Println("    watch:", localLog)
	capData, err := capture.Wait(c, hints, localLog, localCap, opt.Wait)
	if err != nil {
		return res, err
	}
	licenseURL := capData.Get("license_url")
	if strings.TrimSpace(capData.Get("pssh")) == "" && strings.TrimSpace(capData.Get("mpd")) != "" {
		if pssh, err := extractHLSPSSH(capData.Get("mpd")); err == nil && pssh != "" {
			capData["pssh"] = pssh
		}
	}

	fmt.Println("[+] Capture:")
	for _, k := range []string{"pid", "mpd", "license_url", "auth", "pssh"} {
		v := capData.Get(k)
		if v == "" {
			continue
		}
		label := k
		if k == "license_url" {
			label = "license"
		}
		fmt.Println("   ", label+":", trim(v, 110))
	}

	key := ""
	if capData.Get("pssh") != "" && licenseURL != "" {
		k, kerr := fetchKey(opt, root, capData, licenseURL)
		if kerr != nil {
			fmt.Fprintf(os.Stderr, "[!] key fetch skipped: %v\n", kerr)
		} else {
			key = k
			fmt.Println("[+] key:", strings.ReplaceAll(key, "\n", " | "))
		}
	} else {
		fmt.Println("[*] Not enough fields for wvkey (need pssh + license_url) — raw capture saved")
	}

	sess := session.Session{
		Channel:    opt.Channel,
		PSSH:       capData.Get("pssh"),
		Auth:       capData.Get("auth"),
		PID:        capData.Get("pid"),
		Key:        key,
		MPD:        capData.Get("mpd"),
		LicenseURL: licenseURL,
	}
	path, err := session.Write(root, "capture", sess)
	if err != nil {
		return res, err
	}
	fmt.Println("[+] session:", path)
	res.Session = sess
	res.Path = path
	res.Key = key
	res.MPD = sess.MPD
	return res, nil
}

// Run executes the full phone capture pipeline for a registered app module.
func Run(opt Options) (res Result, err error) {
	if opt.App == nil {
		return res, fmt.Errorf("app plugin required (pass --app, or omit --app for passive capture)")
	}
	if opt.Client == nil {
		opt.Client = adb.New()
	}
	if opt.Root == "" {
		opt.Root = repo.Root()
	}
	if opt.Wait <= 0 {
		opt.Wait = 180 * time.Second
	}
	c := opt.Client
	a := opt.App
	root := opt.Root

	// Defers run LIFO: close app first, then clear proxy.
	if opt.ClearProxy {
		defer ClearProxyLater(c)
	}
	if opt.CloseApp {
		defer func() {
			fmt.Printf("[*] Closing %s…\n", a.Package())
			c.ForceStop(a.Package())
		}()
	}

	if err := c.EnsureDevice(); err != nil {
		return res, err
	}

	proxyBin, caHash := resolveProxy(a, root)
	hints := a.CaptureHints()
	useTProxy := false
	if tm, ok := a.(app.TransparentMITM); ok && tm.UseTransparentMITM() {
		useTProxy = true
	}

	if useTProxy {
		// Transparent REDIRECT — keeps phone VPN (BBC UK geo) working.
		proxy.ClearHTTPProxy(c)
		remoteDir := "/data/local/tmp/capture/" + a.Package()
		if err := proxy.StartTransparent(c, proxyBin, a.Package(), "8080", remoteDir, false); err != nil {
			return res, err
		}
		defer proxy.StopTransparent(c)
		hints.RemoteCaps = []string{remoteDir + "/cap.json"}
		hints.RemoteLog = remoteDir + "/appproxy.log"
		fmt.Println("[*] Transparent MITM (VPN OK) — package", a.Package())
	} else {
		if err := proxy.PushAndStart(c, proxyBin, hints.Score.ProxyArgs()...); err != nil {
			return res, err
		}
		if err := proxy.EnsureHTTPProxy(c, ""); err != nil {
			return res, err
		}
		proxy.ReinjectCA(c, caHash)
	}

	if err := a.Launch(c); err != nil {
		fmt.Fprintf(os.Stderr, "[!] launch: %v\n", err)
	}
	if !opt.NoAutoPlay {
		if err := a.AutoPlay(c, opt.Channel); err != nil {
			fmt.Fprintf(os.Stderr, "[!] auto-play failed: %v\n", err)
			fmt.Println("[*] Open the channel on the phone manually and start playback…")
		}
	} else {
		fmt.Println("[*] Open the channel on the phone and wait for playback…")
	}

	_ = os.MkdirAll(filepath.Join(root, ".cache"), 0o755)
	serialTag := c.Serial
	if serialTag == "" {
		serialTag = "default"
	}
	localLog := filepath.Join(root, ".cache", "appproxy-"+sanitize(serialTag)+".log")
	localCap := filepath.Join(root, ".cache", "appproxy-"+sanitize(serialTag)+"_cap.json")
	stop := make(chan struct{})
	go capture.MirrorLogLoop(c, hints.RemoteLog, localLog, stop)
	defer close(stop)

	need := hints.Require
	if len(need) == 0 {
		need = []string{"manifest"}
	}
	fmt.Printf("[*] Waiting for %s…\n", strings.Join(need, " + "))
	fmt.Println("    watch:", localLog)
	capData, err := capture.Wait(c, hints, localLog, localCap, opt.Wait)
	if err != nil {
		return res, err
	}
	licenseURL := first(capData.Get("license_url"), a.LicenseURL())
	keyMode := strings.ToLower(strings.TrimSpace(opt.KeyMode))
	if keyMode == "" || keyMode == "auto" {
		keyMode = strings.ToLower(strings.TrimSpace(a.KeyMode()))
	}
	// Brightcove/HLS: proxy often captures license + master URL but not PSSH.
	// Skip for clear/MPD-only apps (KeyMode none).
	if keyMode != "none" && keyMode != "clear" {
		if strings.TrimSpace(capData.Get("pssh")) == "" && strings.TrimSpace(capData.Get("mpd")) != "" {
			if pssh, err := extractHLSPSSH(capData.Get("mpd")); err == nil && pssh != "" {
				capData["pssh"] = pssh
			} else if err != nil {
				fmt.Fprintf(os.Stderr, "[!] HLS PSSH extract: %v\n", err)
			}
		}
	}

	mpdURL := strings.TrimSpace(capData.Get("mpd"))
	// Console: always surface the MPD we will use (BBC clear streams only need this).
	fmt.Println()
	fmt.Println("========== MPD ==========")
	if mpdURL != "" {
		fmt.Println(mpdURL)
	} else {
		fmt.Println("(none)")
	}
	fmt.Println("=========================")
	fmt.Println()
	fmt.Println("[+] Capture:")
	for _, k := range []string{"pid", "mpd", "license_url", "auth", "pssh"} {
		v := capData.Get(k)
		if k == "license_url" {
			v = first(v, licenseURL)
		}
		if v == "" {
			continue
		}
		label := k
		if k == "license_url" {
			label = "license"
		}
		fmt.Println("   ", label+":", trim(v, 120))
	}

	key := ""
	if keyMode == "none" || keyMode == "clear" {
		fmt.Println("[*] Key mode none — skipping wvkey (clear / MPD-only capture)")
	} else {
		var kerr error
		key, kerr = fetchKey(opt, root, capData, licenseURL)
		if kerr != nil {
			return res, kerr
		}
		fmt.Println("[+] key:", strings.ReplaceAll(key, "\n", " | "))
	}

	sess := session.Session{
		Channel:    opt.Channel,
		PSSH:       capData.Get("pssh"),
		Auth:       capData.Get("auth"),
		PID:        capData.Get("pid"),
		Key:        key,
		MPD:        mpdURL,
		LicenseURL: licenseURL,
	}
	path, err := session.Write(root, a.Name(), sess)
	if err != nil {
		return res, err
	}
	fmt.Println("[+] session:", path)
	if mpdURL != "" {
		fmt.Println("[+] mpd:", mpdURL)
	}
	res.Session = sess
	res.Path = path
	res.Key = key
	res.MPD = sess.MPD
	return res, nil
}

// ClearProxyLater clears http_proxy (exported for agent defer helpers).
func ClearProxyLater(c *adb.Client) {
	proxy.ClearHTTPProxy(c)
}

func resolveProxy(a app.App, root string) (bin, caHash string) {
	bin = filepath.Join(root, "bin", "proxy-android-arm64")
	caHash = DefaultCAHash
	if a != nil {
		if p := a.ProxyBin(); p != "" {
			bin = p
		}
		if h := a.CAHash(); h != "" {
			caHash = h
		}
	}
	if _, err := os.Stat(bin); err != nil {
		for _, p := range []string{
			filepath.Join(root, "bin", "proxy-android-arm64"),
			filepath.Join(root, "apps", "proxy", "proxy-android-arm64"),
		} {
			if st, e := os.Stat(p); e == nil && !st.IsDir() {
				bin = p
				break
			}
		}
	}
	return bin, caHash
}

func fetchKey(opt Options, root string, cap capture.Data, licenseURL string) (string, error) {
	if strings.TrimSpace(opt.WVD) == "" {
		return "", fmt.Errorf("wvkey requires --wvd (path to .wvd device file)")
	}
	py := opt.Python
	if py == "" {
		py = filepath.Join(root, ".venv", "Scripts", "python.exe")
		if _, err := os.Stat(py); err != nil {
			py = filepath.Join(root, ".venv", "bin", "python")
		}
	}
	wopt := wvkey.Options{
		Python:     py,
		Script:     filepath.Join(root, "apps", "wvkey", "wvkey.py"),
		WVD:        opt.WVD,
		PSSH:       cap.Get("pssh"),
		Auth:       cap.Get("auth"),
		PID:        cap.Get("pid"),
		LicenseURL: licenseURL,
		UserAgent:  opt.UserAgent,
	}
	mode := strings.ToLower(strings.TrimSpace(opt.KeyMode))
	if mode == "" || mode == "auto" {
		mode = "modulardrm"
		if opt.App != nil {
			if m := strings.ToLower(strings.TrimSpace(opt.App.KeyMode())); m != "" {
				mode = m
			}
		}
	}
	switch mode {
	case "raw":
		// Brightcove returns multiple CONTENT keys; NRE needs all of them.
		keys, err := wvkey.FetchRawAll(wopt)
		if err != nil {
			return "", err
		}
		return strings.Join(keys, "\n"), nil
	default:
		return wvkey.Fetch(wopt)
	}
}

func first(vals ...string) string {
	for _, v := range vals {
		if strings.TrimSpace(v) != "" {
			return v
		}
	}
	return ""
}

var (
	reWVKeyURI  = regexp.MustCompile(`(?is)KEYFORMAT="urn:uuid:edef8ba9-79d6-4ace-a3c8-27dcd51d21ed".*?URI="data:text/plain;base64,([A-Za-z0-9+/=]+)"`)
	reWVKeyURI2 = regexp.MustCompile(`(?is)URI="data:text/plain;base64,([A-Za-z0-9+/=]+)".*?KEYFORMAT="urn:uuid:edef8ba9-79d6-4ace-a3c8-27dcd51d21ed"`)
)

func extractHLSPSSH(masterURL string) (string, error) {
	resp, err := http.Get(masterURL)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return "", fmt.Errorf("HTTP %d", resp.StatusCode)
	}
	body, err := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
	if err != nil {
		return "", err
	}
	s := string(body)
	if m := reWVKeyURI.FindStringSubmatch(s); len(m) == 2 {
		return m[1], nil
	}
	if m := reWVKeyURI2.FindStringSubmatch(s); len(m) == 2 {
		return m[1], nil
	}
	return "", fmt.Errorf("no Widevine PSSH in HLS master")
}

func trim(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n]
}

func sanitize(s string) string {
	s = strings.Map(func(r rune) rune {
		switch {
		case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9':
			return r
		default:
			return '_'
		}
	}, s)
	return s
}