// Package cluster implements peer-to-peer distribution of the frpc binary and
// an in-memory node registry for the cluster. Binary exchange is preferred
// between cluster members; the external GitHub URL is only a fallback. A new
// member joining asks its neighbors to transmit the binary automatically.
package cluster

import (
	"archive/tar"
	"bytes"
	"compress/gzip"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"sort"
	"strings"
	"time"
)

// PeerTarget identifies a cluster node that can serve binaries.
type PeerTarget struct {
	Addr string // host:port of the peer's frpc binary endpoint
	User string
	Pass string
}

// NodeInfo is exchanged during discovery: each node reports which frpc
// versions it has cached.
type NodeInfo struct {
	ID      string   `json:"id"`
	Addr    string   `json:"addr"`
	Version string   `json:"version,omitempty"`
	Cache   []string `json:"cache,omitempty"`
	SeenAt  int64    `json:"seenAt"`
}

// Registry tracks peers and the locally cached frpc binaries.
type Registry struct {
	BinDir string

	mu      chan struct{} // simple mutex
	peers   map[string]*PeerTarget
	nodeInf map[string]*NodeInfo
}

// NewRegistry builds a Registry rooted at binDir (where frpc-<version> dirs live).
func NewRegistry(binDir string) *Registry {
	return &Registry{
		BinDir:  binDir,
		mu:      make(chan struct{}, 1),
		peers:   map[string]*PeerTarget{},
		nodeInf: map[string]*NodeInfo{},
	}
}

func (r *Registry) lock()   { r.mu <- struct{}{} }
func (r *Registry) unlock() { <-r.mu }

// platformName mirrors install.platformName (kept local to avoid import cycle).
func platformName() string {
	arch := runtime.GOARCH
	if arch == "x86_64" {
		arch = "amd64"
	}
	return runtime.GOOS + "_" + arch
}

// CachedVersions lists frpc versions present under binDir as frpc-<version>.
func (r *Registry) CachedVersions() []string {
	entries, err := os.ReadDir(r.BinDir)
	if err != nil {
		return nil
	}
	var out []string
	for _, e := range entries {
		if !e.IsDir() || !strings.HasPrefix(e.Name(), "frpc-") {
			continue
		}
		out = append(out, strings.TrimPrefix(e.Name(), "frpc-"))
	}
	sort.Strings(out)
	return out
}

// HasVersion reports whether a version is both cached and has a runnable frpc.
func (r *Registry) HasVersion(version string) bool {
	version = strings.TrimPrefix(version, "v")
	p := filepath.Join(r.BinDir, "frpc-"+version, "frpc")
	info, err := os.Stat(p)
	return err == nil && !info.IsDir() && (info.Mode()&0o111) != 0
}

// BinaryPath returns the local file path for a cached version.
func (r *Registry) BinaryPath(version string) string {
	return filepath.Join(r.BinDir, "frpc-"+strings.TrimPrefix(version, "v"), "frpc")
}

// Provide serves a cached binary to a peer. Returns the file path and a
// SHA-256 checksum for verification.
func (r *Registry) Provide(version string) (path, checksum string, err error) {
	version = strings.TrimPrefix(version, "v")
	p := r.BinaryPath(version)
	info, err := os.Stat(p)
	if err != nil || info.IsDir() {
		return "", "", fmt.Errorf("version %s not cached locally", version)
	}
	h := sha256.New()
	f, err := os.Open(p)
	if err != nil {
		return "", "", err
	}
	_, copyErr := io.Copy(h, f)
	f.Close()
	if copyErr != nil {
		return "", "", copyErr
	}
	return p, hex.EncodeToString(h.Sum(nil)), nil
}

// RegisterPeer remembers a reachable peer so binary exchange can prefer it.
func (r *Registry) RegisterPeer(id, addr, user, pass string) {
	r.lock()
	defer r.unlock()
	r.peers[id] = &PeerTarget{Addr: addr, User: user, Pass: pass}
}

// UpdateNodeInfo stores a discovered node's cached versions.
func (r *Registry) UpdateNodeInfo(ni *NodeInfo) {
	r.lock()
	defer r.unlock()
	ni.SeenAt = time.Now().Unix()
	r.nodeInf[ni.ID] = &NodeInfo{ID: ni.ID, Addr: ni.Addr, Version: ni.Version, Cache: append([]string(nil), ni.Cache...), SeenAt: ni.SeenAt}
}

// PeerList returns the currently registered peers (deep copy).
func (r *Registry) PeerList() []*PeerTarget {
	r.lock()
	defer r.unlock()
	out := make([]*PeerTarget, 0, len(r.peers))
	for _, p := range r.peers {
		out = append(out, &PeerTarget{Addr: p.Addr, User: p.User, Pass: p.Pass})
	}
	return out
}

// NodeList returns the discovered nodes (deep copies) including cached versions.
func (r *Registry) NodeList() []NodeInfo {
	r.lock()
	defer r.unlock()
	out := make([]NodeInfo, 0, len(r.nodeInf))
	for _, n := range r.nodeInf {
		out = append(out, NodeInfo{
			ID: n.ID, Addr: n.Addr, Version: n.Version,
			Cache: append([]string(nil), n.Cache...), SeenAt: n.SeenAt,
		})
	}
	return out
}

// fetchFromPeer downloads the frpc binary for a version from one peer over
// HTTP GET /frpc/{version}. Basic Auth is used; the body is capped and its
// SHA-256 must match the peer-provided checksum (defense against corruption).
func fetchFromPeer(ctx context.Context, peer *PeerTarget, version, wantChecksum string) ([]byte, error) {
	url := fmt.Sprintf("http://%s/frpc/%s", peer.Addr, strings.TrimPrefix(version, "v"))
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	if peer.User != "" || peer.Pass != "" {
		req.SetBasicAuth(peer.User, peer.Pass)
	}
	cli := &http.Client{Timeout: 30 * time.Second}
	resp, err := cli.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		_, _ = io.Copy(io.Discard, resp.Body)
		return nil, fmt.Errorf("peer GET %s -> HTTP %d", url, resp.StatusCode)
	}
	data, err := io.ReadAll(io.LimitReader(resp.Body, 256<<20))
	if err != nil {
		return nil, err
	}
	if wantChecksum != "" {
		sum := sha256.Sum256(data)
		got := hex.EncodeToString(sum[:])
		if !strings.EqualFold(got, wantChecksum) {
			return nil, fmt.Errorf("peer binary checksum mismatch: got %s want %s", got, wantChecksum)
		}
	}
	return data, nil
}

// githubURL builds the official GitHub release download URL for a version.
func githubURL(version string) string {
	ver := strings.TrimPrefix(version, "v")
	return fmt.Sprintf("https://github.com/fatedier/frp/releases/download/v%s/frp_%s_%s.tar.gz", ver, ver, platformName())
}

// fetchFromURL downloads the frpc *release tarball* from GitHub and extracts
// the frpc binary (same shape as install.Install). Returns extracted bytes.
func fetchFromURL(ctx context.Context, version string) ([]byte, error) {
	url := githubURL(version)
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	cli := &http.Client{Timeout: 120 * time.Second}
	resp, err := cli.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		_, _ = io.Copy(io.Discard, resp.Body)
		return nil, fmt.Errorf("URL %s -> HTTP %d", url, resp.StatusCode)
	}
	tarGz, err := io.ReadAll(io.LimitReader(resp.Body, 256<<20))
	if err != nil {
		return nil, err
	}
	return extractFrpcBytes(tarGz)
}

// EnsureVersion makes sure a given frpc version is cached locally. It follows
// the distribution priority: local cache -> cluster peers -> external URL.
func (r *Registry) EnsureVersion(ctx context.Context, version string) (string, error) {
	version = strings.TrimPrefix(version, "v")
	if r.HasVersion(version) {
		return r.BinaryPath(version), nil
	}

	// 1) Try peers first (cluster-internal exchange is preferred).
	for _, peer := range r.PeerList() {
		data, err := fetchFromPeer(ctx, peer, version, "")
		if err != nil {
			continue
		}
		if err := r.storeBinary(version, data); err != nil {
			continue
		}
		if r.HasVersion(version) {
			return r.BinaryPath(version), nil
		}
	}

	// 2) Fall back to the external GitHub URL.
	if data, err := fetchFromURL(ctx, version); err == nil {
		if err := r.storeBinary(version, data); err == nil && r.HasVersion(version) {
			return r.BinaryPath(version), nil
		}
	}

	return "", fmt.Errorf("unable to obtain frpc %s from cache, peers, or URL", version)
}

// storeBinary persists freshly fetched frpc bytes under binDir/frpc-<version>
// and verifies the binary runs with --version.
func (r *Registry) storeBinary(version string, data []byte) error {
	if err := os.MkdirAll(r.BinDir, 0o755); err != nil {
		return err
	}
	version = strings.TrimPrefix(version, "v")
	dir := filepath.Join(r.BinDir, "frpc-"+version)
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return err
	}
	bin := filepath.Join(dir, "frpc")
	if err := os.WriteFile(bin, data, 0o755); err != nil {
		return err
	}
	if out, err := exec.Command(bin, "--version").CombinedOutput(); err != nil {
		_ = os.RemoveAll(dir)
		return fmt.Errorf("peer-provided frpc failed --version: %v (%s)", err, bytes.TrimSpace(out))
	}
	return nil
}

// extractFrpcBytes extracts the frpc executable bytes from a GitHub release
// tarball (frp_<v>_<platform>.tar.gz). Only the top-level frpc is taken.
func extractFrpcBytes(tarGz []byte) ([]byte, error) {
	gzr, err := gzip.NewReader(bytes.NewReader(tarGz))
	if err != nil {
		return nil, err
	}
	defer gzr.Close()
	tr := tar.NewReader(gzr)
	for {
		hdr, err := tr.Next()
		if err == io.EOF {
			return nil, fmt.Errorf("frpc not found in archive")
		}
		if err != nil {
			return nil, err
		}
		name := filepath.Base(hdr.Name)
		if hdr.Typeflag == tar.TypeReg && name == "frpc" {
			data, err := io.ReadAll(io.LimitReader(tr, 256<<20))
			if err != nil {
				return nil, err
			}
			return data, nil
		}
	}
}