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"
	"os"
	"os/exec"
	"os/signal"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"
	"time"

	"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.")
}

// ── 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)
		}
		priv, err := ecdh.X25519().NewPrivateKey(raw)
		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)
	}
	priv, err := ecdh.X25519().NewPrivateKey(rawPriv)
	if err != nil {
		fatalf("keygen: %v", err)
	}
	if err := os.MkdirAll(filepath.Dir(keyFile), 0700); err != nil {
		fatalf("mkdir keydir: %v", err)
	}
	if err := os.WriteFile(keyFile, []byte(hex.EncodeToString(rawPriv)+"\n"), 0600); err != nil {
		fatalf("write privkey: %v", err)
	}
	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"`
}

type wsSelfAddress struct {
	Type    string `json:"type"`
	Address string `json:"address"`
}

func printSelfAddress(ws *websocket.Conn) {
	req := map[string]string{"type": "selfAddress"}
	if err := websocket.JSON.Send(ws, req); err != nil {
		fmt.Fprintf(os.Stderr, "selfAddress send: %v\n", err)
		return
	}
	var resp wsSelfAddress
	if err := websocket.JSON.Receive(ws, &resp); err != nil {
		fmt.Fprintf(os.Stderr, "selfAddress recv: %v\n", err)
		return
	}
	fmt.Printf("Nym inbox address: %s\n\n", resp.Address)
}

// ── 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)
		}

		dataB64, ok := extractMessageB64(frame)
		if !ok {
			continue
		}

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

		saveSubmission(outDir, plaintext)
	}
}

// extractMessageB64 pulls the base64 payload string out of a nym-client
// "received" frame. It tolerates the two encodings nym-client uses across
// versions: text/JSON (the common case) and the native binary protocol.
func extractMessageB64(frame []byte) (string, bool) {
	// Text/JSON frame: {"type":"received","message":"<base64>","senderTag":...}
	if frame[0] == '{' {
		var msg wsReceived
		if err := json.Unmarshal(frame, &msg); err != nil || msg.Type != "received" {
			return "", false
		}
		// message is normally a JSON string; tolerate the legacy
		// {"data":"<base64>"} object shape just in case.
		var s string
		if err := json.Unmarshal(msg.Message, &s); err == nil {
			return s, true
		}
		var obj struct {
			Data string `json:"data"`
		}
		if err := json.Unmarshal(msg.Message, &obj); err == nil && obj.Data != "" {
			return obj.Data, true
		}
		return "", false
	}

	// Binary frame (nym native binary protocol):
	//   0x01 (Received tag) | senderTag flag(1) | [senderTag 16 bytes] | message
	// The message bytes are exactly what the sender transmitted — for nymdrop
	// that is the base64 ASCII string produced by the server.
	if frame[0] == 0x01 {
		buf := frame[1:]
		if len(buf) < 1 {
			return "", false
		}
		hasTag := buf[0]
		buf = buf[1:]
		if hasTag == 1 {
			if len(buf) < 16 {
				return "", false
			}
			buf = buf[16:]
		}
		if len(buf) == 0 {
			return "", false
		}
		return string(buf), true
	}

	return "", false
}

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

// decryptMessage decrypts a base64-encoded 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(dataB64 string, privKey *ecdh.PrivateKey) (string, error) {
	raw, err := base64.StdEncoding.DecodeString(dataB64)
	if err != nil {
		return "", fmt.Errorf("base64: %w", err)
	}

	// 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)
	}
	defer zeroBytes(sharedSecret)

	// HKDF-SHA-256
	aesKey := make([]byte, 32)
	hkdfR := hkdf.New(sha256.New, sharedSecret, make([]byte, 32), []byte("nymdrop-v1"))
	if _, err := io.ReadFull(hkdfR, aesKey); err != nil {
		return "", fmt.Errorf("hkdf: %w", err)
	}
	defer zeroBytes(aesKey)

	// AES-GCM-256 decrypt
	block, err := aes.NewCipher(aesKey)
	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 zeroBytes(b []byte) {
	for i := range b {
		b[i] = 0
	}
}

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