package main

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"sync"
	"time"

	"github.com/awnumar/memguard"
	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/nacl/box"
)

const (
	VERSION           = "0.5.0"
	MAX_MESSAGE_SIZE  = 1024 * 1024 // 1MB
	SESSION_TIMEOUT   = 24 * time.Hour
	MAX_SESSIONS      = 1000
	RATE_LIMIT_WINDOW = time.Minute
	MAX_REQUESTS      = 60
)

// Session represents a user session with protected keys and plaintext
type Session struct {
	ID           string
	PublicKey    *memguard.LockedBuffer // NaCl box public key
	PrivateKey   *memguard.LockedBuffer // NaCl box private key
	SharedSecret *memguard.LockedBuffer // Derived shared secret
	Nonce        *memguard.LockedBuffer // Session nonce
	CreatedAt    time.Time
	LastActivity time.Time
	PairedWith   string
	mutex        sync.RWMutex
}

// SessionManager handles all sessions with rate limiting
type SessionManager struct {
	sessions      map[string]*Session
	mutex         sync.RWMutex
	rateLimiter   map[string]*RateLimit
	limiterMutex  sync.RWMutex
}

type RateLimit struct {
	requests  int
	resetTime time.Time
}

// Request/Response types
type StartSessionResponse struct {
	SessionID string `json:"session_id"`
	Status    string `json:"status"`
	Version   string `json:"version"`
}

type PairSessionRequest struct {
	SessionID      string `json:"session_id"`
	BuddySessionID string `json:"buddy_session_id"`
}

type PairSessionResponse struct {
	Status  string `json:"status"`
	Message string `json:"message"`
}

type EncryptRequest struct {
	SessionID string `json:"session_id"`
	Message   string `json:"message"`
}

type EncryptResponse struct {
	EncryptedMessage string `json:"encrypted_message"`
	Status           string `json:"status"`
}

type DecryptRequest struct {
	SessionID        string `json:"session_id"`
	EncryptedMessage string `json:"encrypted_message"`
}

type DecryptResponse struct {
	DecryptedMessage string `json:"decrypted_message"`
	Status           string `json:"status"`
}

type EndSessionRequest struct {
	SessionID string `json:"session_id"`
}

type EndSessionResponse struct {
	Status  string `json:"status"`
	Message string `json:"message"`
}

var sessionManager *SessionManager

func init() {
	sessionManager = &SessionManager{
		sessions:    make(map[string]*Session),
		rateLimiter: make(map[string]*RateLimit),
	}
}

// Rate limiting
func (sm *SessionManager) checkRateLimit(ip string) bool {
	sm.limiterMutex.Lock()
	defer sm.limiterMutex.Unlock()

	now := time.Now()
	limit, exists := sm.rateLimiter[ip]

	if !exists || now.After(limit.resetTime) {
		sm.rateLimiter[ip] = &RateLimit{
			requests:  1,
			resetTime: now.Add(RATE_LIMIT_WINDOW),
		}
		return true
	}

	if limit.requests >= MAX_REQUESTS {
		return false
	}

	limit.requests++
	return true
}

// Generate cryptographically secure session ID
func generateSessionID(publicKey []byte, nonce []byte) string {
	hash := sha256.New()
	hash.Write(publicKey)
	hash.Write(nonce)
	return hex.EncodeToString(hash.Sum(nil))
}

// Session cleanup goroutine
func (sm *SessionManager) startCleanupRoutine() {
	ticker := time.NewTicker(1 * time.Hour)
	go func() {
		for range ticker.C {
			sm.cleanupExpiredSessions()
		}
	}()
}

func (sm *SessionManager) cleanupExpiredSessions() {
	sm.mutex.Lock()
	defer sm.mutex.Unlock()

	now := time.Now()
	for id, session := range sm.sessions {
		if now.Sub(session.LastActivity) > SESSION_TIMEOUT {
			session.destroy()
			delete(sm.sessions, id)
			log.Printf("Expired session cleaned up: %s", id)
		}
	}
}

// Destroy session - securely wipe all sensitive data
func (s *Session) destroy() {
	s.mutex.Lock()
	defer s.mutex.Unlock()

	if s.PrivateKey != nil {
		s.PrivateKey.Destroy()
	}
	if s.PublicKey != nil {
		s.PublicKey.Destroy()
	}
	if s.SharedSecret != nil {
		s.SharedSecret.Destroy()
	}
	if s.Nonce != nil {
		s.Nonce.Destroy()
	}
}

// Middleware for CORS with restrictions
func corsMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		origin := r.Header.Get("Origin")

		// Only allow same-origin or localhost for development
		// In production, set this to your specific domain
		if origin == "" || origin == "http://localhost:8080" || origin == "https://yourdomain.com" {
			w.Header().Set("Access-Control-Allow-Origin", origin)
		}

		w.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type")
		w.Header().Set("Access-Control-Max-Age", "3600")

		if r.Method == "OPTIONS" {
			w.WriteHeader(http.StatusOK)
			return
		}

		next(w, r)
	}
}

// Middleware for rate limiting
func rateLimitMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		ip := r.RemoteAddr

		if !sessionManager.checkRateLimit(ip) {
			http.Error(w, "Rate limit exceeded", http.StatusTooManyRequests)
			return
		}

		next(w, r)
	}
}

// Middleware for request size limiting
func sizeLimitMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		r.Body = http.MaxBytesReader(w, r.Body, MAX_MESSAGE_SIZE)
		next(w, r)
	}
}

// Start session handler
func startSessionHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")

	sessionManager.mutex.Lock()
	if len(sessionManager.sessions) >= MAX_SESSIONS {
		sessionManager.mutex.Unlock()
		http.Error(w, "Maximum sessions reached", http.StatusServiceUnavailable)
		return
	}
	sessionManager.mutex.Unlock()

	// Generate NaCl box key pair (Curve25519)
	publicKey, privateKey, err := box.GenerateKey(rand.Reader)
	if err != nil {
		log.Printf("Key generation failed: %v", err)
		http.Error(w, "Key generation failed", http.StatusInternalServerError)
		return
	}

	// Generate random nonce
	nonce := make([]byte, 24)
	if _, err := rand.Read(nonce); err != nil {
		log.Printf("Nonce generation failed: %v", err)
		http.Error(w, "Nonce generation failed", http.StatusInternalServerError)
		return
	}

	// Generate session ID from public key and nonce
	sessionID := generateSessionID(publicKey[:], nonce)

	// Store keys in memguard protected memory
	session := &Session{
		ID:           sessionID,
		PublicKey:    memguard.NewBufferFromBytes(publicKey[:]),
		PrivateKey:   memguard.NewBufferFromBytes(privateKey[:]),
		Nonce:        memguard.NewBufferFromBytes(nonce),
		CreatedAt:    time.Now(),
		LastActivity: time.Now(),
	}

	sessionManager.mutex.Lock()
	sessionManager.sessions[sessionID] = session
	sessionManager.mutex.Unlock()

	log.Printf("Session started: %s", sessionID)

	resp := StartSessionResponse{
		SessionID: sessionID,
		Status:    "success",
		Version:   VERSION,
	}

	json.NewEncoder(w).Encode(resp)
}

// Pair session handler
func pairSessionHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")

	var req PairSessionRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Invalid request", http.StatusBadRequest)
		return
	}

	if req.SessionID == "" || req.BuddySessionID == "" {
		http.Error(w, "Missing session IDs", http.StatusBadRequest)
		return
	}

	sessionManager.mutex.RLock()
	session, exists1 := sessionManager.sessions[req.SessionID]
	buddySession, exists2 := sessionManager.sessions[req.BuddySessionID]
	sessionManager.mutex.RUnlock()

	if !exists1 || !exists2 {
		http.Error(w, "Session not found", http.StatusNotFound)
		return
	}

	// Update activity timestamps
	session.mutex.Lock()
	session.LastActivity = time.Now()
	session.PairedWith = req.BuddySessionID
	session.mutex.Unlock()

	buddySession.mutex.Lock()
	buddySession.LastActivity = time.Now()
	buddySession.PairedWith = req.SessionID
	buddySession.mutex.Unlock()

	// Compute shared secret using NaCl box
	var sharedSecret [32]byte
	var myPrivate [32]byte
	var theirPublic [32]byte

	copy(myPrivate[:], session.PrivateKey.Bytes())
	copy(theirPublic[:], buddySession.PublicKey.Bytes())

	box.Precompute(&sharedSecret, &theirPublic, &myPrivate)

	// Store shared secret in memguard
	session.mutex.Lock()
	session.SharedSecret = memguard.NewBufferFromBytes(sharedSecret[:])
	session.mutex.Unlock()

	buddySession.mutex.Lock()
	buddySession.SharedSecret = memguard.NewBufferFromBytes(sharedSecret[:])
	buddySession.mutex.Unlock()

	// Zero out stack variables
	for i := range myPrivate {
		myPrivate[i] = 0
	}
	for i := range theirPublic {
		theirPublic[i] = 0
	}
	for i := range sharedSecret {
		sharedSecret[i] = 0
	}

	log.Printf("Sessions paired: %s <-> %s", req.SessionID, req.BuddySessionID)

	resp := PairSessionResponse{
		Status:  "success",
		Message: "Sessions paired successfully",
	}

	json.NewEncoder(w).Encode(resp)
}

// Encrypt handler
func encryptHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")

	var req EncryptRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Invalid request", http.StatusBadRequest)
		return
	}

	if req.SessionID == "" || req.Message == "" {
		http.Error(w, "Missing required fields", http.StatusBadRequest)
		return
	}

	sessionManager.mutex.RLock()
	session, exists := sessionManager.sessions[req.SessionID]
	sessionManager.mutex.RUnlock()

	if !exists {
		http.Error(w, "Session not found", http.StatusNotFound)
		return
	}

	session.mutex.Lock()
	defer session.mutex.Unlock()

	if session.SharedSecret == nil {
		http.Error(w, "Session not paired", http.StatusBadRequest)
		return
	}

	session.LastActivity = time.Now()

	// Store plaintext in memguard protected buffer
	plaintextBuffer := memguard.NewBufferFromBytes([]byte(req.Message))
	defer plaintextBuffer.Destroy()

	// Create AEAD cipher with XChaCha20-Poly1305
	aead, err := chacha20poly1305.NewX(session.SharedSecret.Bytes())
	if err != nil {
		log.Printf("Cipher creation failed: %v", err)
		http.Error(w, "Encryption failed", http.StatusInternalServerError)
		return
	}

	// Generate random nonce
	nonce := make([]byte, aead.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		log.Printf("Nonce generation failed: %v", err)
		http.Error(w, "Encryption failed", http.StatusInternalServerError)
		return
	}

	// Encrypt the message
	ciphertext := aead.Seal(nonce, nonce, plaintextBuffer.Bytes(), nil)

	// Encode to hex for transmission
	encryptedHex := hex.EncodeToString(ciphertext)

	resp := EncryptResponse{
		EncryptedMessage: encryptedHex,
		Status:           "success",
	}

	json.NewEncoder(w).Encode(resp)
}

// Decrypt handler
func decryptHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")

	var req DecryptRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Invalid request", http.StatusBadRequest)
		return
	}

	if req.SessionID == "" || req.EncryptedMessage == "" {
		http.Error(w, "Missing required fields", http.StatusBadRequest)
		return
	}

	sessionManager.mutex.RLock()
	session, exists := sessionManager.sessions[req.SessionID]
	sessionManager.mutex.RUnlock()

	if !exists {
		http.Error(w, "Session not found", http.StatusNotFound)
		return
	}

	session.mutex.Lock()
	defer session.mutex.Unlock()

	if session.SharedSecret == nil {
		http.Error(w, "Session not paired", http.StatusBadRequest)
		return
	}

	session.LastActivity = time.Now()

	// Decode from hex
	ciphertext, err := hex.DecodeString(req.EncryptedMessage)
	if err != nil {
		http.Error(w, "Invalid ciphertext format", http.StatusBadRequest)
		return
	}

	// Create AEAD cipher
	aead, err := chacha20poly1305.NewX(session.SharedSecret.Bytes())
	if err != nil {
		log.Printf("Cipher creation failed: %v", err)
		http.Error(w, "Decryption failed", http.StatusInternalServerError)
		return
	}

	// Extract nonce and ciphertext
	if len(ciphertext) < aead.NonceSize() {
		http.Error(w, "Invalid ciphertext", http.StatusBadRequest)
		return
	}

	nonce := ciphertext[:aead.NonceSize()]
	ciphertext = ciphertext[aead.NonceSize():]

	// Decrypt the message
	plaintext, err := aead.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		log.Printf("Decryption failed: %v", err)
		http.Error(w, "Decryption failed - authentication failed", http.StatusBadRequest)
		return
	}

	// Store plaintext in memguard protected buffer
	plaintextBuffer := memguard.NewBufferFromBytes(plaintext)
	defer plaintextBuffer.Destroy()

	// Zero out plaintext slice
	for i := range plaintext {
		plaintext[i] = 0
	}

	resp := DecryptResponse{
		DecryptedMessage: string(plaintextBuffer.Bytes()),
		Status:           "success",
	}

	json.NewEncoder(w).Encode(resp)
}

// End session handler
func endSessionHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")

	var req EndSessionRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Invalid request", http.StatusBadRequest)
		return
	}

	if req.SessionID == "" {
		http.Error(w, "Missing session ID", http.StatusBadRequest)
		return
	}

	sessionManager.mutex.Lock()
	defer sessionManager.mutex.Unlock()

	session, exists := sessionManager.sessions[req.SessionID]
	if !exists {
		http.Error(w, "Session not found", http.StatusNotFound)
		return
	}

	// Destroy all protected memory
	session.destroy()

	// Remove from sessions map
	delete(sessionManager.sessions, req.SessionID)

	log.Printf("Session terminated: %s", req.SessionID)

	resp := EndSessionResponse{
		Status:  "success",
		Message: "Session terminated and all keys destroyed",
	}

	json.NewEncoder(w).Encode(resp)
}

// Cleanup all sessions on shutdown
func cleanupAllSessions() {
	sessionManager.mutex.Lock()
	defer sessionManager.mutex.Unlock()

	log.Println("Cleaning up all sessions...")

	for id, session := range sessionManager.sessions {
		session.destroy()
		log.Printf("Destroyed session: %s", id)
	}

	sessionManager.sessions = make(map[string]*Session)
	log.Println("All sessions cleaned up")
}

func main() {
	// Initialize memguard
	memguard.CatchInterrupt()
	defer memguard.Purge()

	// Start session cleanup routine
	sessionManager.startCleanupRoutine()

	// Setup HTTP routes with middleware
	http.HandleFunc("/api/start_session", rateLimitMiddleware(corsMiddleware(sizeLimitMiddleware(startSessionHandler))))
	http.HandleFunc("/api/pair_session", rateLimitMiddleware(corsMiddleware(sizeLimitMiddleware(pairSessionHandler))))
	http.HandleFunc("/api/encrypt", rateLimitMiddleware(corsMiddleware(sizeLimitMiddleware(encryptHandler))))
	http.HandleFunc("/api/decrypt", rateLimitMiddleware(corsMiddleware(sizeLimitMiddleware(decryptHandler))))
	http.HandleFunc("/api/end_session", rateLimitMiddleware(corsMiddleware(sizeLimitMiddleware(endSessionHandler))))

	// Serve static files
	http.Handle("/", http.FileServer(http.Dir("./")))

	// Cleanup on exit
	defer cleanupAllSessions()

	port := ":8080"
	log.Printf("NoFuture-Memguard-PQ v%s starting on port %s", VERSION, port)
	log.Printf("Memguard protection active - all keys and plaintext secured")
	log.Println("WARNING: Use HTTPS in production (configure reverse proxy)")

	if err := http.ListenAndServe(port, nil); err != nil {
		log.Fatal("Server failed to start:", err)
	}
}