package main

import (
	"bytes"
	"compress/zlib"
	"context"
	"crypto/hmac"
	"crypto/sha1"
	"encoding/base64"
	"encoding/json"
	"io"
	"net"
	"net/http"
	"strings"
	"time"
)

type userCtxKey struct{}

func withUser(ctx context.Context, u *sessionUser) context.Context {
	return context.WithValue(ctx, userCtxKey{}, u)
}

// currentUser returns the authenticated user for the request, or nil (e.g.
// internal-trust loopback requests carry no user identity).
func currentUser(r *http.Request) *sessionUser {
	u, _ := r.Context().Value(userCtxKey{}).(*sessionUser)
	return u
}

// requireAdmin writes 403 and returns false unless the request is an admin
// (main.py _require_admin).
func requireAdmin(w http.ResponseWriter, r *http.Request) bool {
	if u := currentUser(r); u != nil && u.IsAdmin {
		return true
	}
	writeJSON(w, http.StatusForbidden, map[string]any{"detail": "Admin access required"})
	return false
}

// Session-cookie verification compatible with the Python service's
// itsdangerous URLSafeTimedSerializer(SECRET_KEY) (itsdangerous 2.2):
//   - HMAC-SHA1 signature
//   - django-concat key derivation: sha1(salt + b"signer" + secret_key)
//   - salt "itsdangerous", separator ".", Unix-epoch timestamp
//   - payload = urlsafe-base64(no pad) of compact JSON {"u":username,"a":is_admin},
//     optionally zlib-compressed with a leading "." marker
// Reusing the same SECRET_KEY means a login on the Python service authenticates
// on the Go service during the parallel run.

const sessionMaxAge = 30 * 24 * 3600 // SESSION_MAX_AGE seconds (30 days)

type sessionUser struct {
	Username string
	IsAdmin  bool
}

func deriveSignerKey(secretKey string) []byte {
	h := sha1.New()
	h.Write([]byte("itsdangerous")) // salt
	h.Write([]byte("signer"))
	h.Write([]byte(secretKey))
	return h.Sum(nil)
}

// verifySessionCookie returns the user encoded in a valid, unexpired token, or
// nil. Constant-time signature comparison; never partially trusts a bad token.
func verifySessionCookie(secretKey, token string) *sessionUser {
	if secretKey == "" || token == "" {
		return nil
	}
	// signature is everything after the final separator.
	i := strings.LastIndexByte(token, '.')
	if i <= 0 {
		return nil
	}
	signed := token[:i] // payload + "." + timestamp
	sig, err := base64.RawURLEncoding.DecodeString(token[i+1:])
	if err != nil {
		return nil
	}
	mac := hmac.New(sha1.New, deriveSignerKey(secretKey))
	mac.Write([]byte(signed))
	if !hmac.Equal(sig, mac.Sum(nil)) {
		return nil
	}

	// timestamp is after the second-to-last separator; payload precedes it
	// (the payload may itself start with "." when zlib-compressed).
	j := strings.LastIndexByte(signed, '.')
	if j < 0 {
		return nil
	}
	tsBytes, err := base64.RawURLEncoding.DecodeString(signed[j+1:])
	if err != nil {
		return nil
	}
	var ts int64
	for _, b := range tsBytes {
		ts = ts<<8 | int64(b)
	}
	if time.Now().Unix()-ts > sessionMaxAge {
		return nil // expired
	}

	payload, err := decodeItsdangerousPayload(signed[:j])
	if err != nil {
		return nil
	}
	var data struct {
		U string `json:"u"`
		A bool   `json:"a"`
	}
	if json.Unmarshal(payload, &data) != nil {
		return nil
	}
	return &sessionUser{Username: data.U, IsAdmin: data.A}
}

// issueSessionCookie produces an itsdangerous URLSafeTimedSerializer token
// compatible with the Python service (so Go-issued cookies verify on Python and
// vice-versa). Inverse of verifySessionCookie.
func issueSessionCookie(secretKey string, u sessionUser) string {
	payload, _ := json.Marshal(struct {
		U string `json:"u"`
		A bool   `json:"a"`
	}{u.Username, u.IsAdmin})
	payloadPart := encodeItsdangerousPayload(payload)
	tsPart := base64.RawURLEncoding.EncodeToString(int64ToBytes(time.Now().Unix()))
	signed := payloadPart + "." + tsPart
	mac := hmac.New(sha1.New, deriveSignerKey(secretKey))
	mac.Write([]byte(signed))
	return signed + "." + base64.RawURLEncoding.EncodeToString(mac.Sum(nil))
}

// encodeItsdangerousPayload mirrors URLSafeSerializerBase.dump_payload: zlib-
// compress only when it actually saves more than one byte (it won't for our tiny
// payload), then urlsafe-base64 (no pad), with a "." marker if compressed.
func encodeItsdangerousPayload(jsonBytes []byte) string {
	var buf bytes.Buffer
	zw := zlib.NewWriter(&buf)
	_, _ = zw.Write(jsonBytes)
	_ = zw.Close()
	if compressed := buf.Bytes(); len(compressed) < len(jsonBytes)-1 {
		return "." + base64.RawURLEncoding.EncodeToString(compressed)
	}
	return base64.RawURLEncoding.EncodeToString(jsonBytes)
}

// int64ToBytes encodes a non-negative int as minimal big-endian bytes, matching
// itsdangerous int_to_bytes (verifySessionCookie reads it back the same way).
func int64ToBytes(n int64) []byte {
	if n == 0 {
		return []byte{0}
	}
	var b []byte
	for n > 0 {
		b = append([]byte{byte(n & 0xff)}, b...)
		n >>= 8
	}
	return b
}

func decodeItsdangerousPayload(p string) ([]byte, error) {
	compressed := strings.HasPrefix(p, ".")
	if compressed {
		p = p[1:]
	}
	raw, err := base64.RawURLEncoding.DecodeString(p)
	if err != nil {
		return nil, err
	}
	if !compressed {
		return raw, nil
	}
	zr, err := zlib.NewReader(bytes.NewReader(raw))
	if err != nil {
		return nil, err
	}
	defer zr.Close()
	return io.ReadAll(zr)
}

// authMiddleware replicates main.py's AuthMiddleware: public paths pass through;
// private-source + no X-Forwarded-For is internal-trust (skip auth); otherwise a
// valid session cookie is required.
func (s *Server) authMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if isPublicPath(r.URL.Path) {
			next.ServeHTTP(w, r)
			return
		}
		// Internal trust: only when the peer is private AND nginx did not add
		// X-Forwarded-For (nginx sets XFF on all proxied internet traffic).
		if r.Header.Get("X-Forwarded-For") == "" && isPrivateAddr(clientIP(r)) {
			next.ServeHTTP(w, r)
			return
		}
		if c, err := r.Cookie("session"); err == nil {
			if u := verifySessionCookie(s.secretKey, c.Value); u != nil {
				next.ServeHTTP(w, r.WithContext(withUser(r.Context(), u)))
				return
			}
		}
		if strings.Contains(r.Header.Get("Accept"), "text/html") {
			http.Redirect(w, r, "/login", http.StatusFound)
			return
		}
		writeJSON(w, http.StatusUnauthorized, map[string]any{"detail": "Not authenticated"})
	})
}

func isPublicPath(p string) bool {
	switch p {
	case "/login", "/logout", "/login.html", "/login-style.css", "/health":
		return true
	}
	// WS endpoints authenticate inside their own handlers.
	return strings.HasPrefix(p, "/icons/") || strings.HasPrefix(p, "/ws/")
}

func clientIP(r *http.Request) string {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}

func isPrivateAddr(ip string) bool {
	a := net.ParseIP(ip)
	if a == nil {
		return false
	}
	return a.IsLoopback() || a.IsPrivate()
}