package main

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdh"
	"crypto/rand"
	"crypto/sha256"

	"golang.org/x/crypto/hkdf"
	"embed"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/fs"
	"nymdrop/internal/nym"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"
	"time"

	"github.com/awnumar/memguard"
	"golang.org/x/net/websocket"
)

//go:embed bin/nym-client-linux-amd64
var nymBin embed.FS

func main() {
	keyFile := flag.String("key", "", "path to X25519 private key file (hex, 32 bytes); generated if missing")
	inboxID := flag.String("id", "nymdrop-inbox", "nym-client identity id")
	wsPort  := flag.Int("ws-port", 1977, "nym-client websocket port")
	outDir  := flag.String("out", "", "directory to save submissions (default: ~/nymdrop-inbox)")
	flag.Parse()

	homeDir, _ := os.UserHomeDir()
	if *outDir == "" {
		*outDir = filepath.Join(homeDir, "nymdrop-inbox")
	}
	if err := os.MkdirAll(*outDir, 0700); err != nil {
		fatalf("mkdir outDir: %v", err)
	}

	privKey := loadOrGenKey(*keyFile, homeDir)

	nymBinPath := extractNymClient()
	defer os.RemoveAll(filepath.Dir(nymBinPath))

	configDir := filepath.Join(homeDir, ".nymdrop-reader")
	initNymClient(nymBinPath, *inboxID, configDir)

	nymCmd := startNymClient(nymBinPath, *inboxID, configDir, *wsPort)
	defer nymCmd.Process.Kill()

	wsURL := fmt.Sprintf("ws://127.0.0.1:%d", *wsPort)
	ws := dialWS(wsURL)
	defer ws.Close()

	// Print own Nym address so operator knows where to point sources.
	printSelfAddress(ws)

	fmt.Printf("nymdrop-reader listening — submissions saved to %s\n", *outDir)
	fmt.Println("Press Ctrl+C to quit.")

	go receiveLoop(ws, privKey, *outDir)

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)
	<-sig
	fmt.Println("\nShutting down.")
	// Safety net: wipes any locked buffer still alive (e.g. one in flight in
	// receiveLoop when the signal arrived) on top of the deferred Destroy
	// calls that cover the normal path.
	memguard.Purge()
}

// ── key management ───────────────────────────────────────────────────────────

func loadOrGenKey(keyFile, homeDir string) *ecdh.PrivateKey {
	if keyFile == "" {
		keyFile = filepath.Join(homeDir, ".nymdrop-reader", "privkey.hex")
	}
	data, err := os.ReadFile(keyFile)
	if err == nil {
		raw, err := hex.DecodeString(strings.TrimSpace(string(data)))
		if err != nil {
			fatalf("privkey decode: %v", err)
		}
		// raw is wiped as soon as it's copied into locked, non-swappable memory;
		// ecdh.NewPrivateKey takes its own copy, so the guarded buffer can be
		// destroyed right after construction.
		keyBuf := memguard.NewBufferFromBytes(raw)
		priv, err := ecdh.X25519().NewPrivateKey(keyBuf.Bytes())
		keyBuf.Destroy()
		if err != nil {
			fatalf("privkey load: %v", err)
		}
		fmt.Printf("Loaded private key from %s\n", keyFile)
		pubHex := hex.EncodeToString(priv.PublicKey().Bytes())
		fmt.Printf("Public key (deploy in nymdrop-server): %s\n", pubHex)
		return priv
	}

	// Generate new key.
	rawPriv := make([]byte, 32)
	if _, err := rand.Read(rawPriv); err != nil {
		fatalf("keygen rand: %v", err)
	}
	keyBuf := memguard.NewBufferFromBytes(rawPriv) // wipes rawPriv on copy
	priv, err := ecdh.X25519().NewPrivateKey(keyBuf.Bytes())
	if err != nil {
		keyBuf.Destroy()
		fatalf("keygen: %v", err)
	}
	if err := os.MkdirAll(filepath.Dir(keyFile), 0700); err != nil {
		keyBuf.Destroy()
		fatalf("mkdir keydir: %v", err)
	}
	if err := os.WriteFile(keyFile, []byte(hex.EncodeToString(keyBuf.Bytes())+"\n"), 0600); err != nil {
		keyBuf.Destroy()
		fatalf("write privkey: %v", err)
	}
	keyBuf.Destroy()
	pubHex := hex.EncodeToString(priv.PublicKey().Bytes())
	fmt.Printf("Generated new private key → %s\n", keyFile)
	fmt.Printf("Public key (deploy in nymdrop-server): %s\n\n", pubHex)
	fmt.Println("  Set NYMDROP_PUBKEY_PLACEHOLDER to the public key above in static/crypto.js")
	fmt.Println("  and rebuild the server before accepting submissions.")
	return priv
}

// ── nym-client lifecycle ─────────────────────────────────────────────────────

func extractNymClient() string {
	tmpDir, err := os.MkdirTemp("", "nymdrop-reader-*")
	if err != nil {
		fatalf("mktemp: %v", err)
	}
	binName := "nym-client"
	if runtime.GOOS == "windows" {
		binName += ".exe"
	}
	binPath := filepath.Join(tmpDir, binName)

	srcName := "bin/nym-client-linux-amd64"
	data, err := fs.ReadFile(nymBin, srcName)
	if err != nil {
		fatalf("embedded nym-client not found: %v", err)
	}
	if err := os.WriteFile(binPath, data, 0700); err != nil {
		fatalf("extract nym-client: %v", err)
	}
	return binPath
}

func initNymClient(binPath, id, configDir string) {
	// --home removed: nym-client v1.1.76+ uses ~/.nym/ fixed path
	cmd := exec.Command(binPath, "init", "--id", id)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	_ = cmd.Run()
}

func startNymClient(binPath, id, configDir string, wsPort int) *exec.Cmd {
	cmd := exec.Command(binPath, "run",
		"--id", id,
		"--port", fmt.Sprintf("%d", wsPort),
	)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Start(); err != nil {
		fatalf("nym-client start: %v", err)
	}
	return cmd
}

// dialWS waits for the nym-client WebSocket to be ready, then returns the
// open connection. Reusing the probe connection avoids the duplicate-connection
// panic in nym-client (which rejects a second WS on the same port).
func dialWS(wsURL string) *websocket.Conn {
	fmt.Print("Connecting to Nym mixnet")
	for i := 0; i < 120; i++ {
		ws, err := websocket.Dial(wsURL, "", "http://localhost/")
		if err == nil {
			fmt.Println(" ready.")
			return ws
		}
		fmt.Print(".")
		time.Sleep(500 * time.Millisecond)
	}
	fatalf("nym-client websocket did not start in time")
	return nil
}

// ── websocket message types ───────────────────────────────────────────────────

type wsReceived struct {
	Type      string          `json:"type"`
	Message   json.RawMessage `json:"message"`
	SenderTag string          `json:"senderTag"`
}

// printSelfAddress queries the inbox address using nym-client's binary
// protocol rather than JSON. This is deliberate, not cosmetic: nym-client
// picks text-vs-binary for every later push ("received" events) based on
// the format of the *last* request it saw on this connection (see
// clients/native/src/websocket/handler.rs, ReceivedResponseType). Submission
// payloads are raw binary crypto material with no base64 wrapping (see
// Client.Send in internal/nym) — if this connection stayed in JSON/text
// mode, nym-client would run String::from_utf8_lossy over that binary data
// before handing it to us, silently corrupting it. Querying the self
// address in binary form here locks the connection into binary mode for
// the lifetime of receiveLoop below.
func printSelfAddress(ws *websocket.Conn) {
	if err := websocket.Message.Send(ws, nym.EncodeSelfAddressRequest()); err != nil {
		fmt.Fprintf(os.Stderr, "selfAddress send: %v\n", err)
		return
	}
	var frame []byte
	if err := websocket.Message.Receive(ws, &frame); err != nil {
		fmt.Fprintf(os.Stderr, "selfAddress recv: %v\n", err)
		return
	}
	addr, err := nym.DecodeSelfAddressResponse(frame)
	if err != nil {
		fmt.Fprintf(os.Stderr, "selfAddress decode: %v\n", err)
		return
	}
	fmt.Printf("Nym inbox address: %s\n\n", addr)
}

// ── receive loop ──────────────────────────────────────────────────────────────

func receiveLoop(ws *websocket.Conn, privKey *ecdh.PrivateKey, outDir string) {
	debug := os.Getenv("NYMDROP_DEBUG") != ""
	for {
		// Read the raw frame regardless of opcode. websocket.Message.Receive
		// copies the payload of both text and binary frames into the []byte,
		// whereas websocket.JSON.Receive silently drops anything that is not a
		// well-formed JSON text frame — which is exactly how earlier received
		// messages were being lost.
		var frame []byte
		if err := websocket.Message.Receive(ws, &frame); err != nil {
			if err == io.EOF || strings.Contains(err.Error(), "use of closed network connection") {
				return
			}
			fmt.Fprintf(os.Stderr, "ws recv: %v\n", err)
			time.Sleep(500 * time.Millisecond)
			continue
		}
		if len(frame) == 0 {
			continue
		}
		if debug {
			preview := frame
			if len(preview) > 200 {
				preview = preview[:200]
			}
			fmt.Fprintf(os.Stderr, "[debug] frame: %d bytes, lead=0x%02x: %q\n",
				len(frame), frame[0], preview)
		}

		raw, ok := extractMessage(frame)
		if !ok {
			continue
		}

		plaintext, err := decryptMessage(raw, privKey)
		if err != nil {
			fmt.Fprintf(os.Stderr, "decrypt failed: %v\n", err)
			continue
		}

		saveSubmission(outDir, plaintext)
	}
}

// extractMessage pulls the raw submission payload out of a nym-client
// "received" frame, tolerating both encodings nym-client can use: binary
// (what this connection actually runs in, see printSelfAddress) and
// text/JSON (kept as a defensive fallback only — should never fire here).
func extractMessage(frame []byte) ([]byte, bool) {
	// Binary frame (nym native binary protocol), the expected/only path:
	//   0x01 (Received tag) | has_sender_tag(1) | [sender_tag 16B] | msg_len(8, BE) | msg
	// nym.DecodeReceived returns the message bytes exactly as the sender
	// transmitted them — for nymdrop that is the raw, unencoded packet
	// (ephemeral pubkey || iv || AES-GCM ciphertext), see internal/nym.Client.Send.
	if frame[0] == nym.RespReceived {
		return nym.DecodeReceived(frame)
	}

	// Text/JSON frame: {"type":"received","message":"<base64>","senderTag":...}
	// Only reachable if this connection somehow reverted to text mode; the
	// message field is base64 in that protocol (JSON strings can't carry
	// arbitrary binary), so it must be decoded here to match the binary path.
	if frame[0] == '{' {
		var msg wsReceived
		if err := json.Unmarshal(frame, &msg); err != nil || msg.Type != "received" {
			return nil, false
		}
		var s string
		if err := json.Unmarshal(msg.Message, &s); err != nil {
			return nil, false
		}
		raw, err := base64.StdEncoding.DecodeString(s)
		if err != nil {
			return nil, false
		}
		return raw, true
	}

	return nil, false
}

// ── decryption ────────────────────────────────────────────────────────────────

// decryptMessage decrypts a raw packet from the relay.
//
// Wire format from relay:
//   nonce_hex(32 ASCII chars) + ":" + ephemeral_x25519_pub(32) + iv(12) + ciphertext+tag
//
// The nonce prefix is stripped; it exists only to make identical submissions
// look different on the Nym wire. It is not used in decryption.
func decryptMessage(raw []byte, privKey *ecdh.PrivateKey) (string, error) {
	// Strip relay nonce prefix: 32 hex chars + ':'
	const noncePrefix = 32 + 1 // "deadbeef...:" = 33 bytes
	if len(raw) < noncePrefix+32+12+1 {
		return "", fmt.Errorf("packet too short (%d bytes)", len(raw))
	}
	raw = raw[noncePrefix:]

	// Extract ephemeral public key (32 bytes X25519)
	ephPubBytes := raw[:32]
	iv          := raw[32:44]
	ciphertext  := raw[44:]

	ephPub, err := ecdh.X25519().NewPublicKey(ephPubBytes)
	if err != nil {
		return "", fmt.Errorf("ephemeral pubkey: %w", err)
	}

	// ECDH
	sharedSecret, err := privKey.ECDH(ephPub)
	if err != nil {
		return "", fmt.Errorf("ecdh: %w", err)
	}
	// Locked, non-swappable memory for the two secrets derived per submission;
	// Destroy zeroes and unlocks instead of relying on a plain overwrite loop.
	sharedBuf := memguard.NewBufferFromBytes(sharedSecret)
	defer sharedBuf.Destroy()

	// HKDF-SHA-256
	aesKeyRaw := make([]byte, 32)
	hkdfR := hkdf.New(sha256.New, sharedBuf.Bytes(), make([]byte, 32), []byte("nymdrop-v1"))
	if _, err := io.ReadFull(hkdfR, aesKeyRaw); err != nil {
		return "", fmt.Errorf("hkdf: %w", err)
	}
	aesBuf := memguard.NewBufferFromBytes(aesKeyRaw)
	defer aesBuf.Destroy()

	// AES-GCM-256 decrypt
	block, err := aes.NewCipher(aesBuf.Bytes())
	if err != nil {
		return "", fmt.Errorf("aes: %w", err)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", fmt.Errorf("gcm: %w", err)
	}
	plaintext, err := gcm.Open(nil, iv, ciphertext, nil)
	if err != nil {
		return "", fmt.Errorf("gcm open: %w", err)
	}

	return string(plaintext), nil
}

// ── storage ───────────────────────────────────────────────────────────────────

func saveSubmission(outDir, plaintext string) {
	ts := time.Now().UTC().Format("20060102-150405")
	path := filepath.Join(outDir, ts+".txt")

	// Avoid collision on rapid submissions.
	for i := 1; ; i++ {
		if _, err := os.Stat(path); os.IsNotExist(err) {
			break
		}
		path = filepath.Join(outDir, fmt.Sprintf("%s-%d.txt", ts, i))
	}

	// Separate embedded files from text body.
	body := plaintext
	fileSection := ""
	if idx := strings.Index(plaintext, "\n---FILE:"); idx != -1 {
		body = plaintext[:idx]
		fileSection = plaintext[idx+1:]
	}

	if err := os.WriteFile(path, []byte(body), 0600); err != nil {
		fmt.Fprintf(os.Stderr, "save submission: %v\n", err)
		return
	}
	fmt.Printf("\n[%s] New submission → %s\n", ts, path)
	preview := body
	if len(preview) > 200 {
		preview = preview[:200] + "..."
	}
	fmt.Printf("  %s\n", strings.ReplaceAll(strings.TrimSpace(preview), "\n", "\n  "))

	// Save attached files.
	if fileSection != "" {
		saveAttachments(outDir, ts, fileSection)
	}
}

func saveAttachments(outDir, ts, fileSection string) {
	lines := strings.Split(fileSection, "\n")
	var currentName string
	var currentData strings.Builder

	flush := func() {
		if currentName == "" {
			return
		}
		raw, err := base64.StdEncoding.DecodeString(strings.TrimSpace(currentData.String()))
		if err != nil {
			fmt.Fprintf(os.Stderr, "  attachment decode %q: %v\n", currentName, err)
			return
		}
		safeName := filepath.Base(filepath.Clean(currentName))
		attPath := filepath.Join(outDir, ts+"-"+safeName)
		if err := os.WriteFile(attPath, raw, 0600); err != nil {
			fmt.Fprintf(os.Stderr, "  save attachment %q: %v\n", currentName, err)
			return
		}
		fmt.Printf("  attachment → %s (%d bytes)\n", attPath, len(raw))
		currentName = ""
		currentData.Reset()
	}

	for _, line := range lines {
		if strings.HasPrefix(line, "---FILE:") && strings.HasSuffix(line, "---") {
			flush()
			currentName = strings.TrimSuffix(strings.TrimPrefix(line, "---FILE:"), "---")
		} else if currentName != "" {
			currentData.WriteString(line)
		}
	}
	flush()
}

// ── helpers ───────────────────────────────────────────────────────────────────

func fatalf(format string, args ...any) {
	fmt.Fprintf(os.Stderr, "nymdrop-reader: "+format+"\n", args...)
	os.Exit(1)
}