package auth

import (
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"
)

const (
	sessionTTL       = 8 * time.Hour
	cookieName        = "ar_session"
	nonceExpiry       = 2 * time.Minute
	maxNonces         = 256
	rotateInterval    = 10 * time.Second
	rotateGrace       = 10 * time.Second
	defaultUsername   = "admin"
	defaultPassword   = "admin"
	passwordHashInner = 10000
)

type Session struct {
	CurrentToken  string
	PrevToken     string
	PrevRotatedAt time.Time
	IP            string
	CreatedAt     time.Time
	LastSeen      time.Time
	LastRotate    time.Time
}

type Challenge struct {
	Nonce     string
	CreatedAt time.Time
}

type Store struct {
	mu       sync.Mutex
	sessions map[string]*Session
	nonces   map[string]*Challenge

	muConfig sync.RWMutex
	username string
	password string
	apiKey   string
}

var Global = &Store{
	sessions: make(map[string]*Session),
	nonces:   make(map[string]*Challenge),
}

func (s *Store) Init(username, passwordHash, apiKey string) {
	s.muConfig.Lock()
	defer s.muConfig.Unlock()
	if username != "" {
		s.username = username
	} else {
		s.username = defaultUsername
	}
	if passwordHash != "" {
		s.password = passwordHash
	} else {
		s.password = HashPassword(defaultPassword)
	}
	s.apiKey = apiKey
}

func (s *Store) IsDefaultPassword() bool {
	s.muConfig.RLock()
	defer s.muConfig.RUnlock()
	return s.password == HashPassword(defaultPassword)
}

func (s *Store) GetCredentials() (username, apiKey string, isDefault bool) {
	s.muConfig.RLock()
	defer s.muConfig.RUnlock()
	return s.username, s.apiKey, s.password == HashPassword(defaultPassword)
}

func HashPassword(password string) string {
	salt := "nano-router-salt-v1"
	for i := 0; i < passwordHashInner; i++ {
		data := salt + password + strconv.Itoa(i)
		sum := sha256.Sum256([]byte(data))
		salt = hex.EncodeToString(sum[:16])
	}
	final := sha256.Sum256([]byte(salt + password))
	return hex.EncodeToString(final[:])
}

func generateToken() string {
	b := make([]byte, 32)
	rand.Read(b)
	return hex.EncodeToString(b)
}

func (s *Store) CreateChallenge() string {
	nonce := generateToken()
	s.mu.Lock()
	defer s.mu.Unlock()
	if len(s.nonces) > maxNonces {
		now := time.Now()
		for k, v := range s.nonces {
			if now.Sub(v.CreatedAt) > nonceExpiry {
				delete(s.nonces, k)
			}
		}
	}
	s.nonces[nonce] = &Challenge{Nonce: nonce, CreatedAt: time.Now()}
	return nonce
}

func (s *Store) ValidateChallenge(nonce string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	ch, ok := s.nonces[nonce]
	if !ok {
		return false
	}
	delete(s.nonces, nonce)
	return time.Since(ch.CreatedAt) < nonceExpiry
}

func ComputeResponse(nonce, passwordHash string) string {
	h := sha256.New()
	h.Write([]byte(nonce))
	h.Write([]byte(":"))
	h.Write([]byte(passwordHash))
	return hex.EncodeToString(h.Sum(nil))
}

func (s *Store) Login(nonce, response, clientIP string) (sessionID string, ok bool) {
	if !s.ValidateChallenge(nonce) {
		return "", false
	}

	s.muConfig.RLock()
	defer s.muConfig.RUnlock()

	expected := ComputeResponse(nonce, s.password)
	if subtle.ConstantTimeCompare([]byte(response), []byte(expected)) != 1 {
		return "", false
	}

	sessionID = generateToken()
	session := &Session{
		CurrentToken: sessionID,
		IP:           clientIP,
		CreatedAt:    time.Now(),
		LastSeen:     time.Now(),
		LastRotate:   time.Now(),
	}

	s.mu.Lock()
	s.sessions[sessionID] = session
	s.mu.Unlock()

	return sessionID, true
}

func (s *Store) ValidateAndRotate(token, clientIP string) (newToken string, valid bool) {
	s.mu.Lock()
	defer s.mu.Unlock()

	session, exists := s.sessions[token]
	if !exists {
		return "", false
	}

	if session.IP != clientIP {
		delete(s.sessions, session.CurrentToken)
		if session.PrevToken != "" {
			delete(s.sessions, session.PrevToken)
		}
		return "", false
	}

	if time.Since(session.LastSeen) > sessionTTL {
		delete(s.sessions, session.CurrentToken)
		if session.PrevToken != "" {
			delete(s.sessions, session.PrevToken)
		}
		return "", false
	}

	isPrevToken := token == session.PrevToken
	if isPrevToken {
		if !session.PrevRotatedAt.IsZero() && time.Since(session.PrevRotatedAt) > rotateGrace {
			delete(s.sessions, session.CurrentToken)
			delete(s.sessions, session.PrevToken)
			return "", false
		}
	}

	session.LastSeen = time.Now()

	if time.Since(session.LastRotate) < rotateInterval {
		return "", true
	}

	newTok := generateToken()

	if session.PrevToken != "" {
		delete(s.sessions, session.PrevToken)
	}

	session.PrevToken = session.CurrentToken
	session.PrevRotatedAt = time.Now()
	session.CurrentToken = newTok
	session.LastRotate = time.Now()

	s.sessions[newTok] = session

	return newTok, true
}

func (s *Store) DestroySession(token string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	session, exists := s.sessions[token]
	if !exists {
		return
	}
	if session.PrevToken != "" {
		delete(s.sessions, session.PrevToken)
	}
	delete(s.sessions, session.CurrentToken)
}

func (s *Store) DestroyAllSessions() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.sessions = make(map[string]*Session)
}

func (s *Store) ChangeCredentials(newUsername, newPasswordHash string) {
	s.muConfig.Lock()
	defer s.muConfig.Unlock()
	if newUsername != "" {
		s.username = newUsername
	}
	if newPasswordHash != "" {
		s.password = newPasswordHash
	}
	s.DestroyAllSessions()
}

func (s *Store) VerifyPassword(passwordHash string) bool {
	s.muConfig.RLock()
	defer s.muConfig.RUnlock()
	return subtle.ConstantTimeCompare([]byte(s.password), []byte(passwordHash)) == 1
}

func (s *Store) UpdateUsername(newUsername string) {
	s.muConfig.Lock()
	defer s.muConfig.Unlock()
	if newUsername != "" {
		s.username = newUsername
	}
}

func (s *Store) UpdatePassword(newPasswordHash string, clientIP string) (sessionID string) {
	s.muConfig.Lock()
	s.password = newPasswordHash
	s.muConfig.Unlock()

	s.DestroyAllSessions()

	return s.createSessionAfterPasswordChange(clientIP)
}

func (s *Store) createSessionAfterPasswordChange(clientIP string) string {
	sessionID := generateToken()
	session := &Session{
		CurrentToken: sessionID,
		IP:           clientIP,
		CreatedAt:    time.Now(),
		LastSeen:     time.Now(),
		LastRotate:   time.Now(),
	}
	s.mu.Lock()
	s.sessions[sessionID] = session
	s.mu.Unlock()
	return sessionID
}

func (s *Store) SetAPIKey(key string) {
	s.muConfig.Lock()
	defer s.muConfig.Unlock()
	s.apiKey = key
}

func (s *Store) ValidateAPIKey(key string) bool {
	s.muConfig.RLock()
	defer s.muConfig.RUnlock()
	if s.apiKey == "" {
		return false
	}
	return subtle.ConstantTimeCompare([]byte(key), []byte(s.apiKey)) == 1
}

func (s *Store) Cleanup() {
	now := time.Now()
	s.mu.Lock()
	seen := make(map[*Session]bool)
	for _, session := range s.sessions {
		if seen[session] {
			continue
		}
		seen[session] = true
		if now.Sub(session.LastSeen) > sessionTTL {
			delete(s.sessions, session.CurrentToken)
			if session.PrevToken != "" {
				delete(s.sessions, session.PrevToken)
			}
			continue
		}
		if session.PrevToken != "" && !session.PrevRotatedAt.IsZero() && now.Sub(session.PrevRotatedAt) > rotateGrace {
			delete(s.sessions, session.PrevToken)
			session.PrevToken = ""
		}
	}
	for k, v := range s.nonces {
		if now.Sub(v.CreatedAt) > nonceExpiry {
			delete(s.nonces, k)
		}
	}
	s.mu.Unlock()
}

func GetClientIP(r *http.Request) string {
	if xff := r.Header.Get("X-Forwarded-For"); xff != "" {
		parts := strings.SplitN(xff, ",", 2)
		return strings.TrimSpace(parts[0])
	}
	if xri := r.Header.Get("X-Real-IP"); xri != "" {
		return xri
	}
	ip := r.RemoteAddr
	if idx := strings.LastIndex(ip, ":"); idx != -1 {
		ip = ip[:idx]
	}
	return ip
}

func SetSessionCookie(w http.ResponseWriter, token string) {
	http.SetCookie(w, &http.Cookie{
		Name:     cookieName,
		Value:    token,
		Path:     "/",
		MaxAge:   int(sessionTTL.Seconds()),
		HttpOnly: true,
		SameSite: http.SameSiteStrictMode,
	})
}

func ClearSessionCookie(w http.ResponseWriter) {
	http.SetCookie(w, &http.Cookie{
		Name:     cookieName,
		Value:    "",
		Path:     "/",
		MaxAge:   -1,
		HttpOnly: true,
		SameSite: http.SameSiteStrictMode,
	})
}

func GetSessionToken(r *http.Request) string {
	c, err := r.Cookie(cookieName)
	if err != nil {
		return ""
	}
	return c.Value
}

func AuthMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		token := GetSessionToken(r)
		clientIP := GetClientIP(r)

		if token != "" {
			if newToken, valid := Global.ValidateAndRotate(token, clientIP); valid {
				if newToken != "" {
					SetSessionCookie(w, newToken)
				}
				next.ServeHTTP(w, r)
				return
			}
		}

		authHeader := r.Header.Get("Authorization")
		if strings.HasPrefix(authHeader, "Bearer ") {
			apiKey := strings.TrimPrefix(authHeader, "Bearer ")
			if Global.ValidateAPIKey(apiKey) {
				next.ServeHTTP(w, r)
				return
			}
		}

		if isHTMLRequest(r) {
			http.Redirect(w, r, "/login.html", http.StatusFound)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusUnauthorized)
		fmt.Fprintf(w, `{"success":false,"error":"unauthorized"}`)
	})
}

func PublicAuthMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/login.html" || r.URL.Path == "/api/auth/challenge" || r.URL.Path == "/api/auth/login" {
			next.ServeHTTP(w, r)
			return
		}

		if strings.HasPrefix(r.URL.Path, "/api/auth/") || r.URL.Path == "/style.css" {
			next.ServeHTTP(w, r)
			return
		}

		token := GetSessionToken(r)
		clientIP := GetClientIP(r)

		if token != "" {
			if newToken, valid := Global.ValidateAndRotate(token, clientIP); valid {
				if newToken != "" {
					SetSessionCookie(w, newToken)
				}
				next.ServeHTTP(w, r)
				return
			}
		}

		authHeader := r.Header.Get("Authorization")
		if strings.HasPrefix(authHeader, "Bearer ") {
			apiKey := strings.TrimPrefix(authHeader, "Bearer ")
			if Global.ValidateAPIKey(apiKey) {
				next.ServeHTTP(w, r)
				return
			}
		}

		if isHTMLRequest(r) {
			http.Redirect(w, r, "/login.html", http.StatusFound)
			return
		}

		if strings.HasPrefix(r.URL.Path, "/api/") {
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusUnauthorized)
			fmt.Fprintf(w, `{"success":false,"error":"unauthorized"}`)
			return
		}

		next.ServeHTTP(w, r)
	})
}

func isHTMLRequest(r *http.Request) bool {
	accept := r.Header.Get("Accept")
	return strings.Contains(accept, "text/html") ||
		strings.HasSuffix(r.URL.Path, ".html") ||
		r.URL.Path == "/" ||
		(!strings.HasPrefix(r.URL.Path, "/api/") && !strings.Contains(accept, "application/json"))
}

func StartCleanup(interval time.Duration) {
	go func() {
		ticker := time.NewTicker(interval)
		defer ticker.Stop()
		for range ticker.C {
			Global.Cleanup()
		}
	}()
}