package httpserver

import (
	"crypto/aes"
	"crypto/cipher"
	"bytes"
	"encoding/base64"
	"encoding/json"
	"io"
	"log"
	"net/http"
	"net/http/httputil"
	"sync"
	"time"

	cr "aesdemo/server/internal/crypto"
	lu "aesdemo/server/internal/logutil"
)

type decryptRequest struct {
	Nonce      string `json:"nonce"`
	Ciphertext string `json:"ciphertext"`
	Alg        string `json:"alg"`
}

type decryptResponse struct {
	OK   bool   `json:"ok"`
	Error string `json:"error,omitempty"`
}

// serverState holds runtime options and a short-TTL cache of recent nonces.
type serverState struct {
    k1        []byte
    debugHTTP bool
    logPlain  bool
    nonceTTL  time.Duration
    mu        sync.Mutex
    recent    map[string]time.Time // key: nonce base64, value: first-seen time
}

const defaultNonceTTL = 2 * time.Minute

func (s *serverState) isReplayAndStore(nonceB64 string) bool {
    now := time.Now()
    cutoff := now.Add(-s.nonceTTL)
    s.mu.Lock()
    defer s.mu.Unlock()
    // prune expired
    for n, t := range s.recent {
        if t.Before(cutoff) {
            delete(s.recent, n)
        }
    }
    if t, ok := s.recent[nonceB64]; ok && t.After(cutoff) {
        return true
    }
    s.recent[nonceB64] = now
    return false
}

func NewMuxWithOptions(k1 []byte, debugHTTP bool, logPlain bool) http.Handler {
	mux := http.NewServeMux()
	s := &serverState{
		k1:        k1,
		debugHTTP: debugHTTP,
		logPlain:  logPlain,
		nonceTTL:  defaultNonceTTL,
		recent:    make(map[string]time.Time),
	}
	log.Printf("[data plane] Anti-replay enabled: nonce_ttl=%s", s.nonceTTL.String())
	mux.HandleFunc("/v1/decrypt", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			w.WriteHeader(http.StatusMethodNotAllowed)
			return
		}
		// Dump full HTTP request (headers + body) only if debug enabled
		if s.debugHTTP && r.Body != nil {
			bodyBytes, _ := io.ReadAll(r.Body)
			_ = r.Body.Close()
			// First dump
			r.Body = io.NopCloser(bytes.NewReader(bodyBytes))
			if dump, err := httputil.DumpRequest(r, true); err == nil {
				log.Printf("[data plane] full request from client:\n%s", string(dump))
			}
			// Restore for JSON decoder
			r.Body = io.NopCloser(bytes.NewReader(bodyBytes))
		}
		var req decryptRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			writeJSON(w, http.StatusBadRequest, decryptResponse{OK: false, Error: "invalid json"})
			return
		}
		if req.Alg != cr.AlgorithmName {
			writeJSON(w, http.StatusBadRequest, decryptResponse{OK: false, Error: "unsupported alg"})
			return
		}
		salt, err := base64.StdEncoding.DecodeString(req.Nonce)
		if err != nil {
			writeJSON(w, http.StatusBadRequest, decryptResponse{OK: false, Error: "bad nonce"})
			return
		}
		ct, err := base64.StdEncoding.DecodeString(req.Ciphertext)
		if err != nil {
			writeJSON(w, http.StatusBadRequest, decryptResponse{OK: false, Error: "bad ciphertext"})
			return
		}
		log.Printf("[data plane] HTTP request: %s %s len=%d content-type=%s", r.Method, r.URL.Path, r.ContentLength, r.Header.Get("Content-Type"))
		log.Printf("[data plane] Envelope: alg=%s nonce_b64_len=%d ct_b64_len=%d", req.Alg, len(req.Nonce), len(req.Ciphertext))
		log.Printf("[data plane] Nonce received: base64=%s len=%d", req.Nonce, len(salt))

		// Anti-replay: reject if nonce seen recently
		if s.isReplayAndStore(req.Nonce) {
			log.Printf("[data plane] Replay detected: nonce_base64=%s (within %s)", req.Nonce, s.nonceTTL.String())
			writeJSON(w, http.StatusConflict, decryptResponse{OK: false, Error: "replay detected"})
			return
		}

		// Derive K2 and GCM nonce, log derivation info
		k2, gcmNonce, err := cr.DeriveKeyAndGcmNonce(s.k1, salt)
		if err != nil {
			log.Printf("[data plane] Derive fail: %v", err)
			writeJSON(w, http.StatusBadRequest, decryptResponse{OK: false, Error: "derive failed"})
			return
		}
		k1fp := lu.FingerprintSHA256(s.k1, 4)
		k2fp := lu.FingerprintSHA256(k2, 4)
		log.Printf("[data plane] Derive: k1_fp=%s hkdf=SHA256 info=\"aes-demo-v1\" => k2_fp=%s gcm_nonce_b64=%s", k1fp, k2fp, base64.StdEncoding.EncodeToString(gcmNonce))

		// Decrypt using derived values
		blk, err := aes.NewCipher(k2)
		if err != nil {
			log.Printf("[data plane] Decrypt fail: %v", err)
			writeJSON(w, http.StatusBadRequest, decryptResponse{OK: false, Error: "cipher init failed"})
			return
		}
		a, err := cipher.NewGCMWithNonceSize(blk, len(gcmNonce))
		if err != nil {
			log.Printf("[data plane] Decrypt fail: %v", err)
			writeJSON(w, http.StatusBadRequest, decryptResponse{OK: false, Error: "gcm init failed"})
			return
		}
		pt, err := a.Open(nil, gcmNonce, ct, nil)
		if err != nil {
			log.Printf("[data plane] Decrypt fail: %v", err)
			writeJSON(w, http.StatusBadRequest, decryptResponse{OK: false, Error: "decrypt failed"})
			return
		}
		ptB64 := base64.StdEncoding.EncodeToString(pt)
		log.Printf("[data plane] Decrypt ok: alg=%s, pt_len=%d, pt_b64_len=%d, pt_str=%q", req.Alg, len(pt), len(ptB64), string(pt))
		writeJSON(w, http.StatusOK, decryptResponse{OK: true})
	})
	return mux
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}