package proxy

import (
	"crypto/md5"
	"crypto/x509"
	"encoding/binary"
	"encoding/pem"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"drmdecryption/adb"
	"drmdecryption/repo"
)

const (
	// Universal on-device paths (still kill legacy rteproxy process names).
	RemoteBin     = "/data/local/tmp/appproxy"
	RemoteCap     = "/data/local/tmp/appproxy_cap.json"
	RemoteLog     = "/data/local/tmp/appproxy.log"
	RemoteTraffic = "/data/local/tmp/appproxy_traffic.jsonl"
	RemoteCACrt   = "/data/local/tmp/rteproxy-ca.crt"
	DefaultCAHash = "6c3578b4"
)

var reWLANIPv4 = regexp.MustCompile(`(?m)^\s*inet\s+(\d{1,3}(?:\.\d{1,3}){3})/`)

// DeviceWLANIPv4 returns the phone's current wlan0 IPv4 address.
func DeviceWLANIPv4(c *adb.Client) (string, error) {
	out := c.Out("shell", "ip", "-f", "inet", "addr", "show", "wlan0")
	if m := reWLANIPv4.FindStringSubmatch(out); len(m) == 2 {
		return m[1], nil
	}
	// Fallbacks used on some OEM builds.
	for _, prop := range []string{"dhcp.wlan0.ipaddress", "dhcp.eth0.ipaddress"} {
		if v := c.Out("shell", "getprop", prop); net.ParseIP(v) != nil && v != "0.0.0.0" {
			return v, nil
		}
	}
	return "", fmt.Errorf("no wlan0 IPv4 (is Wi‑Fi connected?)")
}

// DeviceHTTPProxyAddr returns "<wlan-ip>:port" for the on-device MITM.
// Loopback (127.0.0.1) must not be used: after MITM, some clients rewrite the
// upstream Host to the proxy address, and appproxy then dials 127.0.0.1:443.
func DeviceHTTPProxyAddr(c *adb.Client, port string) (string, error) {
	if strings.TrimSpace(port) == "" {
		port = "8080"
	}
	ip, err := DeviceWLANIPv4(c)
	if err != nil {
		return "", err
	}
	return net.JoinHostPort(ip, port), nil
}

func looksLikeLoopbackProxy(want string) bool {
	host, _, err := net.SplitHostPort(strings.TrimSpace(want))
	if err != nil {
		host = strings.TrimSpace(want)
	}
	host = strings.Trim(host, "[]")
	return host == "127.0.0.1" || host == "localhost" || host == "::1" || host == "0.0.0.0" || host == ""
}

// EnsureHTTPProxy points the device at the on-device mitm.
// Empty or loopback want is rewritten to the phone's WLAN IP:8080 so upstream
// MITM dials keep the real destination host (BBC/RTE/etc.).
func EnsureHTTPProxy(c *adb.Client, want string) error {
	if looksLikeLoopbackProxy(want) {
		port := "8080"
		if hostport := strings.TrimSpace(want); hostport != "" {
			if _, p, err := net.SplitHostPort(hostport); err == nil && p != "" {
				port = p
			}
		}
		addr, err := DeviceHTTPProxyAddr(c, port)
		if err != nil {
			return err
		}
		want = addr
	}
	cur := c.Out("shell", "settings", "get", "global", "http_proxy")
	if cur == want {
		fmt.Println("[+] HTTP proxy already", want)
		return nil
	}
	fmt.Printf("[*] Setting HTTP proxy -> %s (was %q)\n", want, cur)
	_, err := c.Shell("settings", "put", "global", "http_proxy", want)
	got := c.Out("shell", "settings", "get", "global", "http_proxy")
	if got != want {
		return fmt.Errorf("failed to set http_proxy (got %q)", got)
	}
	return err
}

// ClearHTTPProxy disables the global HTTP proxy and stops any leftover mitm.
// Always call this after a capture/agent job so the phone can play apps normally.
func ClearHTTPProxy(c *adb.Client) {
	_, _ = c.Shell("settings", "put", "global", "http_proxy", ":0")
	_, _ = c.Shell("settings", "delete", "global", "http_proxy")
	_, _ = c.Shell("settings", "delete", "global", "global_http_proxy_host")
	_, _ = c.Shell("settings", "delete", "global", "global_http_proxy_port")
	_, _ = c.Shell("settings", "delete", "global", "global_http_proxy_exclusion_list")
	stopProxy(c)
	fmt.Println("[*] HTTP proxy cleared")
}

func stopProxy(c *adb.Client) {
	script := `
for name in appproxy rteproxy; do
  pid=$(pidof $name 2>/dev/null || true)
  if [ -n "$pid" ]; then kill $pid >/dev/null 2>&1 || true; fi
done
sleep 0.5
for name in appproxy rteproxy; do
  pid=$(pidof $name 2>/dev/null || true)
  if [ -n "$pid" ]; then kill -9 $pid >/dev/null 2>&1 || true; fi
done
# Root fallback — some builds ignore non-root kill.
if command -v su >/dev/null 2>&1; then
  su -c 'killall appproxy rteproxy 2>/dev/null; killall -9 appproxy rteproxy 2>/dev/null; true' 2>/dev/null || true
fi
sleep 0.3
(pidof appproxy || pidof rteproxy) >/dev/null 2>&1 && echo STILL || echo STOPPED`
	out, _ := c.Shell("sh", "-c", script)
	if strings.Contains(out, "STILL") {
		fmt.Println("[!] old appproxy still running after stop — port 8080 may stay busy")
	}
}

// extraFlags renders extra device flags, quoting each value.
func extraFlags(args []string) string {
	if len(args) == 0 {
		return ""
	}
	var b strings.Builder
	for _, a := range args {
		b.WriteString(" ")
		if strings.HasPrefix(a, "-") {
			b.WriteString(a)
			continue
		}
		b.WriteString("'" + strings.ReplaceAll(a, "'", "") + "'")
	}
	return b.String()
}

// PushAndStart installs and launches the on-device mitm binary. extraArgs are
// appended to its command line — app modules pass their manifest-scoring hosts
// that way, since the device binary cannot read a module's values file.
func PushAndStart(c *adb.Client, localBin string, extraArgs ...string) error {
	if st, err := os.Stat(localBin); err != nil || st.IsDir() {
		return fmt.Errorf("missing proxy binary %s — build apps/proxy first (proxyctl build)", localBin)
	}
	fmt.Println("[*] Stopping any old appproxy/rteproxy...")
	stopProxy(c)

	fmt.Println("[*] Pushing appproxy...")
	if err := c.Push(localBin, RemoteBin); err != nil {
		return err
	}
	_, _ = c.Shell("chmod", "755", RemoteBin)
	_, _ = c.Shell("rm", "-f", RemoteCap, "/data/local/tmp/rte_cap.json", "/sdcard/Download/rte_cap.json", "/storage/emulated/0/Download/rte_cap.json", RemoteLog, "/data/local/tmp/rteproxy.log")

	starter := "#!/system/bin/sh\n" +
		"exec " + RemoteBin +
		" -listen :8080" +
		" -out " + RemoteCap +
		" -ca-dir /data/local/tmp" +
		" -dns 1.1.1.1,1.0.0.1,8.8.8.8,192.168.1.1" +
		" >>" + RemoteLog + " 2>&1\n"

	tmp := filepath.Join(os.TempDir(), "start_appproxy.sh")
	if err := os.WriteFile(tmp, []byte(starter), 0o755); err != nil {
		return err
	}
	defer os.Remove(tmp)
	if err := c.Push(tmp, "/data/local/tmp/start_appproxy.sh"); err != nil {
		return err
	}
	_, _ = c.Shell("chmod", "755", "/data/local/tmp/start_appproxy.sh")

	// Keep backgrounded after adb exits.
	_, _, _ = c.Run("shell", "setsid /data/local/tmp/start_appproxy.sh </dev/null >/dev/null 2>&1 &")

	var pid, logHead string
	for i := 0; i < 10; i++ {
		time.Sleep(400 * time.Millisecond)
		pid = c.Out("shell", "pidof", "appproxy")
		if pid == "" {
			pid = c.Out("shell", "pidof", "rteproxy")
		}
		logHead = c.Out("shell", "head", "-12", RemoteLog)
		if pid != "" && strings.Contains(logHead, "listening") {
			break
		}
	}
	if logHead != "" {
		fmt.Println(logHead)
	}
	errLog := c.Out("shell", "cat", RemoteLog)
	if strings.Contains(errLog, "address already in use") ||
		(strings.Contains(errLog, "listen ") && strings.Contains(errLog, "bind:")) {
		fmt.Fprintln(os.Stderr, errLog)
		return fmt.Errorf("appproxy failed to bind :8080 (old process still holding the port?)")
	}
	if pid == "" || !strings.Contains(logHead, "listening") {
		if errLog != "" {
			fmt.Fprintln(os.Stderr, errLog)
		}
		return fmt.Errorf("appproxy failed to start")
	}
	fmt.Printf("[+] appproxy pid=%s; capture -> %s\n", pid, RemoteCap)
	return nil
}

// FindLocalBin returns the first existing proxy binary under the repo.
func FindLocalBin(root string) string {
	if root == "" {
		root = repo.Root()
	}
	for _, p := range []string{
		filepath.Join(root, "bin", "proxy-android-arm64"),
		filepath.Join(root, "apps", "proxy", "proxy-android-arm64"),
		filepath.Join(root, "apps", "proxy", "rteproxy-android-arm64"),
		filepath.Join(root, "bin", "rteproxy-android-arm64"),
	} {
		if st, err := os.Stat(p); err == nil && !st.IsDir() {
			return p
		}
	}
	return ""
}

// FindLocalCA returns the first existing MITM CA cert under the repo.
func FindLocalCA(root string) string {
	if root == "" {
		root = repo.Root()
	}
	for _, p := range []string{
		filepath.Join(root, "apps", "proxy", "rteproxy-ca.crt"),
		filepath.Join(root, "data", "proxy-ca.crt"),
	} {
		if st, err := os.Stat(p); err == nil && !st.IsDir() {
			return p
		}
	}
	return ""
}

// AndroidCAHash returns the OpenSSL subject_hash_old used for system CA files.
func AndroidCAHash(certPEM []byte) (string, error) {
	block, _ := pem.Decode(certPEM)
	if block == nil {
		return "", fmt.Errorf("no PEM certificate found")
	}
	cert, err := x509.ParseCertificate(block.Bytes)
	if err != nil {
		return "", err
	}
	sum := md5.Sum(cert.RawSubject)
	n := binary.LittleEndian.Uint32(sum[0:4])
	return fmt.Sprintf("%08x", n), nil
}

// InstallCA pushes the MITM CA onto the device as HASH.0 and optionally Magisk-reinjects it.
// When reinject is true, runs the full Magisk conscrypt bind (needs root / Magisk busybox).
func InstallCA(c *adb.Client, localCA string, reinject bool) (hash string, err error) {
	if localCA == "" {
		localCA = FindLocalCA("")
	}
	if localCA == "" {
		return "", fmt.Errorf("no local CA cert found (expected apps/proxy/rteproxy-ca.crt)")
	}
	pemBytes, err := os.ReadFile(localCA)
	if err != nil {
		return "", err
	}
	hash, err = AndroidCAHash(pemBytes)
	if err != nil {
		return "", err
	}
	fmt.Printf("[*] Installing CA %s (hash %s)\n", localCA, hash)

	if err := c.Push(localCA, RemoteCACrt); err != nil {
		return hash, err
	}
	// Also drop a copy named after the process for clarity.
	_ = c.Push(localCA, "/data/local/tmp/appproxy-ca.crt")

	tmpHash := filepath.Join(os.TempDir(), hash+".0")
	if err := os.WriteFile(tmpHash, pemBytes, 0o644); err != nil {
		return hash, err
	}
	defer os.Remove(tmpHash)
	remoteHash := "/data/local/tmp/" + hash + ".0"
	if err := c.Push(tmpHash, remoteHash); err != nil {
		return hash, err
	}
	_, _ = c.Shell("chmod", "644", remoteHash, RemoteCACrt)
	fmt.Printf("[+] Pushed %s and %s\n", RemoteCACrt, remoteHash)

	if reinject {
		ReinjectCA(c, hash)
	} else {
		fmt.Println("[*] Skipped Magisk reinject (pass -reinject / install-ca --reinject)")
		fmt.Println("    User-CA path: Settings → Security → Install a certificate → CA certificate")
	}
	return hash, nil
}

// Stop stops the on-device mitm process.
func Stop(c *adb.Client) {
	stopProxy(c)
}

const (
	RemoteTProxyScript = "/data/local/tmp/tproxy_iptables.sh"
	RemoteCaptureRoot  = "/data/local/tmp/capture"
)

// StartTransparent pushes appproxy in -transparent mode (no Wi‑Fi http_proxy),
// installs iptables UID REDIRECT for pkg, and writes captures under remoteDir.
func StartTransparent(c *adb.Client, localBin, pkg, port, remoteDir string, reinject bool) error {
	if port == "" {
		port = "8080"
	}
	if remoteDir == "" {
		remoteDir = RemoteCaptureRoot + "/" + pkg
	}
	stopProxy(c)
	_ = stopTransparentRules(c)

	// Ensure no global HTTP proxy — transparent mode must not use one.
	ClearHTTPProxy(c)

	if reinject {
		ReinjectCA(c, DefaultCAHash)
	}

	if st, err := os.Stat(localBin); err != nil || st.IsDir() {
		return fmt.Errorf("missing proxy binary %s", localBin)
	}
	fmt.Println("[*] Pushing appproxy (transparent)...")
	if err := c.Push(localBin, RemoteBin); err != nil {
		return err
	}
	_, _ = c.Shell("chmod", "755", RemoteBin)

	scriptLocal := filepath.Join(repo.Root(), "apps", "proxy", "device", "tproxy_iptables.sh")
	if _, err := os.Stat(scriptLocal); err != nil {
		return fmt.Errorf("missing %s", scriptLocal)
	}
	if err := c.Push(scriptLocal, RemoteTProxyScript); err != nil {
		return err
	}
	_, _ = c.Shell("chmod", "755", RemoteTProxyScript)
	_, _ = c.Shell("mkdir", "-p", remoteDir)
	_, _ = c.Shell("rm", "-f", remoteDir+"/cap.json", remoteDir+"/traffic.jsonl", remoteDir+"/appproxy.log")

	starter := "#!/system/bin/sh\n" +
		"mkdir -p '" + remoteDir + "'\n" +
		"exec " + RemoteBin +
		" -transparent" +
		" -listen :" + port +
		" -mitm-internal 127.0.0.1:18080" +
		" -out-dir '" + remoteDir + "'" +
		" -ca-dir /data/local/tmp" +
		" -dns 1.1.1.1,1.0.0.1,8.8.8.8" +
		" -log-all" +
		" -v" +
		"\n"
	tmp := filepath.Join(os.TempDir(), "start_appproxy_tproxy.sh")
	if err := os.WriteFile(tmp, []byte(starter), 0o755); err != nil {
		return err
	}
	defer os.Remove(tmp)
	if err := c.Push(tmp, "/data/local/tmp/start_appproxy_tproxy.sh"); err != nil {
		return err
	}
	_, _ = c.Shell("chmod", "755", "/data/local/tmp/start_appproxy_tproxy.sh")
	_, _, _ = c.Run("shell", "setsid /data/local/tmp/start_appproxy_tproxy.sh </dev/null >/dev/null 2>&1 &")

	var pid string
	for i := 0; i < 15; i++ {
		time.Sleep(400 * time.Millisecond)
		pid = c.Out("shell", "pidof", "appproxy")
		if pid != "" {
			break
		}
	}
	if pid == "" {
		return fmt.Errorf("appproxy failed to start (transparent)")
	}
	fmt.Printf("[+] appproxy pid=%s; capture → %s\n", pid, remoteDir)

	out, errOut, err := c.Run("shell", "su", "-c", "sh "+RemoteTProxyScript+" start "+pkg+" "+port)
	fmt.Print(out)
	if err != nil {
		return fmt.Errorf("iptables: %v (%s)", err, strings.TrimSpace(errOut+out))
	}
	fmt.Println("[+] iptables REDIRECT installed (UK VPN can stay ON; do not set Wi‑Fi proxy)")
	return nil
}

// StopTransparent removes iptables rules and stops appproxy (does not require Wi‑Fi proxy clear beyond safety).
func StopTransparent(c *adb.Client) {
	_ = stopTransparentRules(c)
	stopProxy(c)
	ClearHTTPProxy(c)
	fmt.Println("[*] transparent capture stopped")
}

func stopTransparentRules(c *adb.Client) error {
	out, errOut, err := c.Run("shell", "su", "-c", "sh "+RemoteTProxyScript+" stop")
	if strings.TrimSpace(out) != "" {
		fmt.Print(out)
	}
	if err != nil && !strings.Contains(errOut+out, "STOPPED") {
		return fmt.Errorf("iptables stop: %v %s", err, errOut)
	}
	return nil
}

// PullDir pulls regular files under remoteDir into destDir.
// Avoids `adb shell sh -c "ls …"` — on Windows that often lists `/` instead of the path.
func PullDir(c *adb.Client, remoteDir, destDir string) error {
	if err := os.MkdirAll(destDir, 0o755); err != nil {
		return err
	}
	remoteDir = strings.TrimRight(strings.TrimSpace(remoteDir), "/")
	// Prefer known capture artifacts; fall back to find -type f.
	// Do not use `adb shell sh -c "ls …"` — on Windows that often lists `/`.
	var names []string
	for _, n := range []string{"cap.json", "traffic.jsonl", "appproxy.log"} {
		if fileExistsOnDevice(c, remoteDir+"/"+n) {
			names = append(names, n)
		}
	}
	if len(names) == 0 {
		list := c.Out("shell", "find", remoteDir, "-maxdepth", "1", "-type", "f")
		for _, line := range strings.Split(list, "\n") {
			line = strings.TrimSpace(line)
			if line == "" {
				continue
			}
			names = append(names, filepath.Base(filepath.Clean(line)))
		}
	}
	if len(names) == 0 {
		return fmt.Errorf("no files in %s", remoteDir)
	}
	seen := map[string]bool{}
	pulled := 0
	for _, name := range names {
		name = strings.TrimSpace(name)
		if name == "" || name == "." || name == ".." || strings.ContainsAny(name, `/\`) || seen[name] {
			continue
		}
		seen[name] = true
		remote := remoteDir + "/" + name
		local := filepath.Join(destDir, name)
		if err := c.Pull(remote, local); err != nil {
			fmt.Printf("[!] pull %s: %v\n", remote, err)
			continue
		}
		fmt.Printf("[+] %s\n", local)
		pulled++
	}
	if pulled == 0 {
		return fmt.Errorf("failed to pull any files from %s", remoteDir)
	}
	return nil
}

func fileExistsOnDevice(c *adb.Client, remote string) bool {
	_, _, err := c.Run("shell", "ls", remote)
	return err == nil
}

// PullCaptures pulls traffic/cap/log into destDir (created if missing).
func PullCaptures(c *adb.Client, destDir string) error {
	if err := os.MkdirAll(destDir, 0o755); err != nil {
		return err
	}
	for _, remote := range []string{RemoteTraffic, RemoteCap, RemoteLog} {
		base := filepath.Base(remote)
		local := filepath.Join(destDir, base)
		if err := c.Pull(remote, local); err != nil {
			fmt.Printf("[!] pull %s: %v\n", remote, err)
			continue
		}
		fmt.Printf("[+] %s\n", local)
	}
	return nil
}

// DiscoverOutDir returns outputs/discover/<stamp> under the repo root.
func DiscoverOutDir(root, stamp string) string {
	if root == "" {
		root = repo.Root()
	}
	if stamp == "" {
		stamp = time.Now().Format("20060102-150405")
	}
	return filepath.Join(root, "outputs", "discover", stamp)
}

// ReinjectCA best-effort Magisk bind of the mitm CA into conscrypt.
func ReinjectCA(c *adb.Client, caHash string) {
	if caHash == "" {
		caHash = DefaultCAHash
	}
	// Bound the per-app nsenter loop — wait on hundreds of PIDs can hang forever.
	script := fmt.Sprintf(`
BB=/data/adb/magisk/busybox
HASH=%s
[ -f /data/local/tmp/$HASH.0 ] || { echo CA_SKIP; exit 0; }
[ -x "$BB" ] || { echo CA_SKIP; exit 0; }
rm -rf /data/local/tmp/cacerts-overlay
mkdir -p /data/local/tmp/cacerts-overlay
cp /apex/com.android.conscrypt/cacerts/* /data/local/tmp/cacerts-overlay/ 2>/dev/null || true
cp /data/local/tmp/$HASH.0 /data/local/tmp/cacerts-overlay/$HASH.0
chmod 644 /data/local/tmp/cacerts-overlay/*
$BB mount -t tmpfs tmpfs /system/etc/security/cacerts 2>/dev/null || true
cp /data/local/tmp/cacerts-overlay/* /system/etc/security/cacerts/ 2>/dev/null || true
chmod 644 /system/etc/security/cacerts/* 2>/dev/null || true
chcon u:object_r:system_file:s0 /system/etc/security/cacerts/* 2>/dev/null || true
$BB mount --bind /system/etc/security/cacerts /apex/com.android.conscrypt/cacerts 2>/dev/null || true
for Z in $(pidof zygote64 2>/dev/null); do
  nsenter --mount=/proc/$Z/ns/mnt -- $BB mount --bind /system/etc/security/cacerts /apex/com.android.conscrypt/cacerts 2>/dev/null || true
done
# Only rebind the target app if set; otherwise skip the full zygote-child sweep (hangs).
PKG_UID_FILE=/data/local/tmp/reinject_target_uid
if [ -f "$PKG_UID_FILE" ]; then
  TUID=$(cat "$PKG_UID_FILE")
  for PID in $(ps -A -o PID=,UID= 2>/dev/null | awk -v u="$TUID" '$2==u {print $1}'); do
    nsenter --mount=/proc/$PID/ns/mnt -- $BB mount --bind /system/etc/security/cacerts /apex/com.android.conscrypt/cacerts 2>/dev/null || true
  done
fi
ls /apex/com.android.conscrypt/cacerts/$HASH.0 2>/dev/null && echo CA_OK || echo CA_SKIP
`, caHash)
	fmt.Println("[*] Re-injecting system CA (Magisk)...")
	tmp := filepath.Join(os.TempDir(), "reinject_ca.sh")
	_ = os.WriteFile(tmp, []byte("#!/system/bin/sh\n"+script), 0o755)
	defer os.Remove(tmp)
	_ = c.Push(tmp, "/data/local/tmp/reinject_ca.sh")
	_, _ = c.Shell("chmod", "755", "/data/local/tmp/reinject_ca.sh")
	// Host-side timeout: su -mm + nsenter has hung on some Magisk builds.
	type runResult struct {
		out string
	}
	ch := make(chan runResult, 1)
	go func() {
		out, _, _ := c.Run("shell", "su", "-mm", "-c", "sh /data/local/tmp/reinject_ca.sh")
		ch <- runResult{out: out}
	}()
	var out string
	select {
	case r := <-ch:
		out = r.out
	case <-time.After(20 * time.Second):
		fmt.Println("[!] CA reinject timed out after 20s — continuing (CA likely already mounted)")
		_, _ = c.Shell("su", "-c", "killall reinject_ca.sh 2>/dev/null; true")
		return
	}
	if strings.Contains(out, "CA_OK") {
		fmt.Println("[+] System CA present in conscrypt")
	} else {
		fmt.Println("[!] CA inject skipped/failed — if TLS errors, re-run Magisk CA mount")
	}
}