// Package reader holds the nymdrop-reader engine: it runs an embedded
// nym-client, receives submissions over its WebSocket, decrypts them, and
// writes them to disk. It is deliberately transport/UI-agnostic — both the
// headless CLI (cmd/nymdrop-reader) and the Fyne GUI drive it through Hooks
// instead of writing to stdout directly, so the two front-ends can share one
// implementation instead of drifting apart.
package reader

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/ecdh"
	"crypto/rand"
	"crypto/sha256"
	"embed"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"io/fs"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"time"

	"golang.org/x/crypto/hkdf"

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

	"nymdrop/internal/nym"
)

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

// Config holds everything needed to start a reader instance. Empty fields
// fall back to the historical defaults (~/.nymdrop-reader, ~/nymdrop-inbox)
// so existing deployments (e.g. the pietro systemd service) keep behaving
// identically when DataDir is left unset.
type Config struct {
	KeyFile string // path to X25519 private key file (hex); generated if missing
	InboxID string // nym-client identity id
	WSPort  int    // nym-client websocket port
	OutDir  string // directory to save decrypted submissions
	// DataDir, if non-empty, makes the whole instance self-contained under
	// this directory: the embedded nym-client's own config/state (normally
	// fixed to ~/.nym by nym-client itself) is redirected here by setting
	// HOME for the child process, and KeyFile/OutDir default inside it too.
	// This is what makes a USB-portable, zero-trace-on-host reader possible
	// without any cooperation from nym-client itself (verified: nym-client
	// reads $HOME like any other Rust `dirs`-crate consumer, nothing more
	// specific is exposed on its CLI).
	DataDir string
	Debug   bool
}

// Hooks lets a front-end observe reader activity without depending on how
// the engine is implemented. Any field left nil is a no-op.
type Hooks struct {
	Log         func(format string, args ...any)
	SelfAddress func(addr string)
	Submission  func(path, preview string)
	Attachment  func(path string, size int)
	Error       func(err error)
}

func (h *Hooks) logf(format string, args ...any) {
	if h.Log != nil {
		h.Log(format, args...)
	}
}

func (h *Hooks) errorf(err error) {
	if h.Error != nil {
		h.Error(err)
	}
}

// Instance is a running reader: an embedded nym-client subprocess plus the
// WebSocket connection to it.
type Instance struct {
	cfg        Config
	hooks      Hooks
	privKey    *ecdh.PrivateKey
	nymCmd     *exec.Cmd
	nymBinDir  string
	ws         *websocket.Conn
	SelfAddr   string
}

// resolveDefaults fills empty Config fields the same way the original
// nymdrop-reader main() did, anchoring everything under DataDir when set.
func resolveDefaults(cfg Config) (Config, error) {
	if cfg.InboxID == "" {
		cfg.InboxID = "nymdrop-inbox"
	}
	if cfg.WSPort == 0 {
		cfg.WSPort = 1977
	}

	base := cfg.DataDir
	if base == "" {
		home, err := os.UserHomeDir()
		if err != nil {
			return cfg, fmt.Errorf("home dir: %w", err)
		}
		base = home
	}

	if cfg.KeyFile == "" {
		cfg.KeyFile = filepath.Join(base, ".nymdrop-reader", "privkey.hex")
	}
	if cfg.OutDir == "" {
		cfg.OutDir = filepath.Join(base, "nymdrop-inbox")
	}
	return cfg, nil
}

// Start loads/generates the decryption key, extracts and launches the
// embedded nym-client, and connects to its WebSocket. It returns once the
// connection is ready and the self-address has been queried; call
// ReceiveLoop afterwards (in its own goroutine if driving a GUI) and Close
// when done.
func Start(cfg Config, hooks Hooks) (*Instance, error) {
	cfg, err := resolveDefaults(cfg)
	if err != nil {
		return nil, err
	}
	if err := os.MkdirAll(cfg.OutDir, 0700); err != nil {
		return nil, fmt.Errorf("mkdir outDir: %w", err)
	}

	privKey, err := loadOrGenKey(cfg.KeyFile, &hooks)
	if err != nil {
		return nil, err
	}

	nymBinPath, nymBinDir, err := extractNymClient()
	if err != nil {
		return nil, err
	}

	// nym-client's own config directory is hardcoded to $HOME/.nym (no CLI
	// flag for it) — redirecting HOME for the child process is the only way
	// to keep a portable run from touching the operator's real home. See
	// Config.DataDir doc comment. Left empty (cfg.DataDir=="") this is a
	// no-op: the child inherits the real environment, identical to before.
	nymHome := cfg.DataDir

	if err := initNymClient(nymBinPath, cfg.InboxID, nymHome); err != nil {
		os.RemoveAll(nymBinDir)
		return nil, err
	}

	nymCmd, err := startNymClient(nymBinPath, cfg.InboxID, cfg.WSPort, nymHome)
	if err != nil {
		os.RemoveAll(nymBinDir)
		return nil, err
	}

	wsURL := fmt.Sprintf("ws://127.0.0.1:%d", cfg.WSPort)
	hooks.logf("Connecting to Nym mixnet")
	ws, err := dialWS(wsURL, &hooks)
	if err != nil {
		nymCmd.Process.Kill()
		os.RemoveAll(nymBinDir)
		return nil, err
	}

	in := &Instance{
		cfg:       cfg,
		hooks:     hooks,
		privKey:   privKey,
		nymCmd:    nymCmd,
		nymBinDir: nymBinDir,
		ws:        ws,
	}

	in.SelfAddr = in.querySelfAddress()
	if in.SelfAddr != "" && hooks.SelfAddress != nil {
		hooks.SelfAddress(in.SelfAddr)
	}

	return in, nil
}

// OutDir returns the directory submissions are written to.
func (in *Instance) OutDir() string { return in.cfg.OutDir }

// Close tears down the nym-client subprocess, its WebSocket, and its
// extracted binary's temp directory.
func (in *Instance) Close() {
	if in.ws != nil {
		in.ws.Close()
	}
	if in.nymCmd != nil && in.nymCmd.Process != nil {
		in.nymCmd.Process.Kill()
	}
	if in.nymBinDir != "" {
		os.RemoveAll(in.nymBinDir)
	}
	memguard.Purge()
}

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

func loadOrGenKey(keyFile string, hooks *Hooks) (*ecdh.PrivateKey, error) {
	data, err := os.ReadFile(keyFile)
	if err == nil {
		raw, err := hex.DecodeString(strings.TrimSpace(string(data)))
		if err != nil {
			return nil, fmt.Errorf("privkey decode: %w", err)
		}
		keyBuf := memguard.NewBufferFromBytes(raw)
		priv, err := ecdh.X25519().NewPrivateKey(keyBuf.Bytes())
		keyBuf.Destroy()
		if err != nil {
			return nil, fmt.Errorf("privkey load: %w", err)
		}
		hooks.logf("Loaded private key from %s", keyFile)
		hooks.logf("Public key (deploy in nymdrop-server): %s", hex.EncodeToString(priv.PublicKey().Bytes()))
		return priv, nil
	}

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

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

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

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

// childEnv returns the environment for the embedded nym-client subprocess,
// overriding HOME (and USERPROFILE, for a future Windows build of
// nym-client) to nymHome when it is non-empty. Verified against the real
// binary: nym-client happily initialises its ~/.nym tree under a redirected
// HOME with no other flag needed.
func childEnv(nymHome string) []string {
	if nymHome == "" {
		return nil // inherit parent's environment unchanged
	}
	env := os.Environ()
	filtered := env[:0]
	for _, kv := range env {
		if strings.HasPrefix(kv, "HOME=") || strings.HasPrefix(kv, "USERPROFILE=") {
			continue
		}
		filtered = append(filtered, kv)
	}
	filtered = append(filtered, "HOME="+nymHome, "USERPROFILE="+nymHome)
	return filtered
}

func initNymClient(binPath, id, nymHome string) error {
	cmd := exec.Command(binPath, "init", "--id", id)
	cmd.Env = childEnv(nymHome)
	cmd.Stdout = io.Discard
	cmd.Stderr = io.Discard
	_ = cmd.Run() // best-effort, same as before: a pre-existing id is fine
	return nil
}

func startNymClient(binPath, id string, wsPort int, nymHome string) (*exec.Cmd, error) {
	cmd := exec.Command(binPath, "run",
		"--id", id,
		"--port", fmt.Sprintf("%d", wsPort),
	)
	cmd.Env = childEnv(nymHome)
	cmd.Stdout = io.Discard
	cmd.Stderr = io.Discard
	if err := cmd.Start(); err != nil {
		return nil, fmt.Errorf("nym-client start: %w", err)
	}
	return cmd, nil
}

func dialWS(wsURL string, hooks *Hooks) (*websocket.Conn, error) {
	for i := 0; i < 120; i++ {
		ws, err := websocket.Dial(wsURL, "", "http://localhost/")
		if err == nil {
			hooks.logf("Nym mixnet connection ready.")
			return ws, nil
		}
		time.Sleep(500 * time.Millisecond)
	}
	return nil, fmt.Errorf("nym-client websocket did not start in time")
}

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

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

// querySelfAddress 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 (in *Instance) querySelfAddress() string {
	if err := websocket.Message.Send(in.ws, nym.EncodeSelfAddressRequest()); err != nil {
		in.hooks.errorf(fmt.Errorf("selfAddress send: %w", err))
		return ""
	}
	var frame []byte
	if err := websocket.Message.Receive(in.ws, &frame); err != nil {
		in.hooks.errorf(fmt.Errorf("selfAddress recv: %w", err))
		return ""
	}
	addr, err := nym.DecodeSelfAddressResponse(frame)
	if err != nil {
		in.hooks.errorf(fmt.Errorf("selfAddress decode: %w", err))
		return ""
	}
	return addr
}

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

// ReceiveLoop blocks, decrypting and saving submissions as they arrive.
// Call it in its own goroutine when driving a GUI; it returns when the
// WebSocket is closed (e.g. by Close).
func (in *Instance) ReceiveLoop() {
	for {
		var frame []byte
		if err := websocket.Message.Receive(in.ws, &frame); err != nil {
			if err == io.EOF || strings.Contains(err.Error(), "use of closed network connection") {
				return
			}
			in.hooks.errorf(fmt.Errorf("ws recv: %w", err))
			time.Sleep(500 * time.Millisecond)
			continue
		}
		if len(frame) == 0 {
			continue
		}
		if in.cfg.Debug {
			preview := frame
			if len(preview) > 200 {
				preview = preview[:200]
			}
			in.hooks.logf("[debug] frame: %d bytes, lead=0x%02x: %q", len(frame), frame[0], preview)
		}

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

		plaintext, err := decryptMessage(raw, in.privKey)
		if err != nil {
			in.hooks.errorf(fmt.Errorf("decrypt failed: %w", err))
			continue
		}

		in.saveSubmission(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 querySelfAddress) and
// text/JSON (kept as a defensive fallback only — should never fire here).
func extractMessage(frame []byte) ([]byte, bool) {
	if frame[0] == nym.RespReceived {
		return nym.DecodeReceived(frame)
	}

	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) {
	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:]

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

	sharedSecret, err := privKey.ECDH(ephPub)
	if err != nil {
		return "", fmt.Errorf("ecdh: %w", err)
	}
	sharedBuf := memguard.NewBufferFromBytes(sharedSecret)
	defer sharedBuf.Destroy()

	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()

	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 (in *Instance) saveSubmission(plaintext string) {
	ts := time.Now().UTC().Format("20060102-150405")
	path := filepath.Join(in.cfg.OutDir, ts+".txt")

	for i := 1; ; i++ {
		if _, err := os.Stat(path); os.IsNotExist(err) {
			break
		}
		path = filepath.Join(in.cfg.OutDir, fmt.Sprintf("%s-%d.txt", ts, i))
	}

	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 {
		in.hooks.errorf(fmt.Errorf("save submission: %w", err))
		return
	}
	preview := body
	if len(preview) > 200 {
		preview = preview[:200] + "..."
	}
	if in.hooks.Submission != nil {
		in.hooks.Submission(path, strings.TrimSpace(preview))
	}

	if fileSection != "" {
		in.saveAttachments(ts, fileSection)
	}
}

func (in *Instance) saveAttachments(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 {
			in.hooks.errorf(fmt.Errorf("attachment decode %q: %w", currentName, err))
			return
		}
		safeName := filepath.Base(filepath.Clean(currentName))
		attPath := filepath.Join(in.cfg.OutDir, ts+"-"+safeName)
		if err := os.WriteFile(attPath, raw, 0600); err != nil {
			in.hooks.errorf(fmt.Errorf("save attachment %q: %w", currentName, err))
			return
		}
		if in.hooks.Attachment != nil {
			in.hooks.Attachment(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()
}