package capture

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"drmdecryption/adb"
)

// Hints tell the waiter which fields / log tags mean a capture is complete.
type Hints struct {
	RemoteCaps []string // paths on device to pull
	RemoteLog  string
	Require    []string // json keys: auth, pid, pssh, mpd
	// BestEffort: return early when any useful field appears; at timeout return
	// whatever was collected (error only if completely empty).
	BestEffort bool
	MPDLogRE   *regexp.Regexp
	// Score ranks candidate manifest URLs. Zero value falls back to
	// DefaultScoreCfg, which has no provider hosts.
	Score ScoreCfg
}

// RemoteCapPaths are the on-device locations the MITM may persist its capture
// JSON to. The rte_cap.json entries are legacy names still written by proxies
// already deployed on phones in the field.
var RemoteCapPaths = []string{
	"/data/local/tmp/appproxy_cap.json",
	"/data/local/tmp/rte_cap.json",
	"/sdcard/Download/rte_cap.json",
	"/storage/emulated/0/Download/rte_cap.json",
}

// RemoteLogPath is where the on-device MITM writes its log.
const RemoteLogPath = "/data/local/tmp/appproxy.log"

// DefaultHints is the device-layout baseline every app starts from. Callers set
// Require (and optionally Score) for their own DRM shape.
func DefaultHints() Hints {
	return Hints{
		RemoteCaps: append([]string{}, RemoteCapPaths...),
		RemoteLog:  RemoteLogPath,
		MPDLogRE:   regexp.MustCompile(`\[MPD\] (https://\S+)`),
		Score:      DefaultScoreCfg(),
	}
}

// DefaultPassiveHints is for a bare capture run: dump whatever the MITM sees.
func DefaultPassiveHints() Hints {
	h := DefaultHints()
	h.BestEffort = true
	return h
}

// score applies the hints' URL scoring, defaulting when unset.
func (h Hints) score(u string) int {
	if len(h.Score.Deny) == 0 && len(h.Score.Hosts) == 0 {
		return DefaultScoreCfg().Score(u)
	}
	return h.Score.Score(u)
}

// Data is the merged capture payload before key fetch.
type Data map[string]string

func (d Data) Get(k string) string { return d[k] }

// Wait pulls device JSON + local mirrored log until required fields exist.
func Wait(c *adb.Client, hints Hints, localLog string, localCap string, timeout time.Duration) (Data, error) {
	deadline := time.Now().Add(timeout)
	lastNote := ""
	for time.Now().Before(deadline) {
		cap := Data{}
		for _, remote := range hints.RemoteCaps {
			_ = os.Remove(localCap)
			if err := c.Pull(remote, localCap); err != nil {
				continue
			}
			raw, err := os.ReadFile(localCap)
			if err != nil {
				continue
			}
			var m map[string]any
			if json.Unmarshal(raw, &m) != nil {
				continue
			}
			for k, v := range m {
				if s, ok := v.(string); ok && s != "" {
					cap[k] = s
				}
			}
			break
		}

		logText := ""
		if b, err := os.ReadFile(localLog); err == nil {
			logText = string(b)
		}
		enrichFromProxyLog(cap, logText, hints)

		if hints.BestEffort {
			if hasUseful(cap, hints) {
				_ = os.WriteFile(localCap, mustJSON(cap), 0o644)
				return cap, nil
			}
		} else {
			missing := false
			for _, k := range hints.Require {
				if strings.TrimSpace(cap[k]) == "" {
					missing = true
					break
				}
			}
			if !missing {
				_ = os.WriteFile(localCap, mustJSON(cap), 0o644)
				return cap, nil
			}
		}

		have := []string{}
		keys := hints.Require
		if hints.BestEffort {
			keys = []string{"mpd", "license_url", "pssh", "auth", "pid"}
		}
		for _, k := range keys {
			if cap[k] != "" {
				have = append(have, k)
			}
		}
		note := "json=" + strings.Join(have, ",")
		if note == "json=" {
			note = "json=none"
		}
		if strings.Contains(logText, "[LIC]") {
			note += " log=LIC"
		}
		if strings.Contains(logText, "[PSSH]") {
			note += " log=PSSH"
		}
		if strings.Contains(logText, "[MPD]") || strings.Contains(logText, "[MAN]") || strings.Contains(logText, "[MEDIA]") || strings.Contains(logText, "[MS]") {
			note += " log=MAN"
		}
		if note != lastNote {
			fmt.Println("   …" + note)
			lastNote = note
		}
		time.Sleep(time.Second)
	}
	if hints.BestEffort {
		// Final pull after timeout — return whatever we got.
		cap := Data{}
		for _, remote := range hints.RemoteCaps {
			_ = os.Remove(localCap)
			if err := c.Pull(remote, localCap); err != nil {
				continue
			}
			raw, err := os.ReadFile(localCap)
			if err != nil {
				continue
			}
			var m map[string]any
			if json.Unmarshal(raw, &m) != nil {
				continue
			}
			for k, v := range m {
				if s, ok := v.(string); ok && s != "" {
					cap[k] = s
				}
			}
			break
		}
		logText := ""
		if b, err := os.ReadFile(localLog); err == nil {
			logText = string(b)
		}
		enrichFromProxyLog(cap, logText, hints)
		if len(cap) > 0 {
			_ = os.WriteFile(localCap, mustJSON(cap), 0o644)
			return cap, nil
		}
		return nil, fmt.Errorf("timed out — no DRM/manifest traffic seen (play something on the phone)")
	}
	return nil, fmt.Errorf("timed out waiting for capture (%v)", hints.Require)
}

func hasUseful(cap Data, hints Hints) bool {
	// Real stream signal only — ignore catalog/EPG URLs parked in "mpd".
	if mpd := strings.TrimSpace(cap["mpd"]); mpd != "" && hints.score(mpd) > 0 {
		return true
	}
	for _, k := range []string{"license_url", "pssh", "auth", "pid"} {
		if strings.TrimSpace(cap[k]) != "" {
			return true
		}
	}
	return false
}

var (
	reLogLIC   = regexp.MustCompile(`\[LIC\]\s+(?:POST|GET)\s+(https://\S+)`)
	reLogMPD   = regexp.MustCompile(`\[MPD\]\s+(https://\S+)`)
	reLogMAN   = regexp.MustCompile(`\[MAN\]\s+GET\s+(https://\S+)`)
	reLogMEDIA = regexp.MustCompile(`\[MEDIA\]\s+(https://\S+)`)
	reLogPSSH = regexp.MustCompile(`\[PSSH\]\s+(\S+)`)
)

func enrichFromProxyLog(cap Data, logText string, hints Hints) {
	if logText == "" {
		return
	}
	bestMPD := ""
	bestScore := 0
	consider := func(u string) {
		u = strings.TrimRight(u, ".,)")
		if u == "" || !strings.HasPrefix(u, "http") {
			return
		}
		sc := hints.score(u)
		// Ignore EPG/schedule/metrics noise (score <= 0).
		if sc <= 0 {
			return
		}
		if sc > bestScore {
			bestScore = sc
			bestMPD = u
		}
	}
	if hints.MPDLogRE != nil {
		for _, m := range hints.MPDLogRE.FindAllStringSubmatch(logText, -1) {
			consider(m[1])
		}
	}
	for _, m := range reLogMPD.FindAllStringSubmatch(logText, -1) {
		consider(m[1])
	}
	for _, m := range reLogMAN.FindAllStringSubmatch(logText, -1) {
		consider(m[1])
	}
	for _, m := range reLogMEDIA.FindAllStringSubmatch(logText, -1) {
		consider(m[1])
	}
	if bestMPD != "" && (cap["mpd"] == "" || hints.score(bestMPD) > hints.score(cap["mpd"])) {
		cap["mpd"] = bestMPD
	}
	// Drop a previously stored non-manifest "mpd" (e.g. schedules feed).
	if cap["mpd"] != "" && hints.score(cap["mpd"]) <= 0 {
		delete(cap, "mpd")
	}
	if cap["license_url"] == "" {
		if ms := reLogLIC.FindAllStringSubmatch(logText, -1); len(ms) > 0 {
			cap["license_url"] = strings.TrimRight(ms[len(ms)-1][1], ".,)")
		}
	}
	if cap["pssh"] == "" {
		if ms := reLogPSSH.FindAllStringSubmatch(logText, -1); len(ms) > 0 {
			cap["pssh"] = ms[len(ms)-1][1]
		}
	}
}

func mustJSON(d Data) []byte {
	b, _ := json.MarshalIndent(map[string]string(d), "", "  ")
	return append(b, '\n')
}

// MirrorLog starts a background `adb exec-out log` equivalent via continuous pull.
// For v1 we periodically pull the remote log file into localLog.
func MirrorLogLoop(c *adb.Client, remote, local string, stop <-chan struct{}) {
	_ = os.MkdirAll(filepath.Dir(local), 0o755)
	for {
		select {
		case <-stop:
			return
		default:
			_ = c.Pull(remote, local)
			time.Sleep(800 * time.Millisecond)
		}
	}
}