package main

import (
	"crypto/subtle"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"runtime/debug"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

const objectPrefix = "/objects/"

type ObjectStore interface {
	Put(key string, body io.Reader, contentType string) (ObjectInfo, error)
	Open(key string) (*ObjectRead, error)
	Stat(key string) (ObjectInfo, error)
	Delete(key string) error
	DiskStatus() (DiskStatus, error)
	List(prefix string, maxKeys int) ([]ObjectInfo, error)
	ListPage(prefix string, after string, maxKeys int) (ListPage, error)
}

type Server struct {
	store         ObjectStore
	cfg           Config
	rateLimiter   *RateLimiter
	metrics       *Metrics
	writeMu       sync.Mutex
	reservedBytes uint64
}

func NewServer(store ObjectStore, cfg Config) http.Handler {
	return NewServerInstance(store, cfg).Handler()
}

func NewServerInstance(store ObjectStore, cfg Config) *Server {
	cfg = cfg.withDefaults()
	return &Server{store: store, cfg: cfg, rateLimiter: NewRateLimiter(cfg.RateLimitPerMinute), metrics: &Metrics{}}
}

func (s *Server) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("/health", s.health)
	mux.HandleFunc("/metrics", s.metricsHandler)
	mux.HandleFunc(objectPrefix, s.object)
	mux.HandleFunc("/", s.s3Object)
	return s.recoverPanics(s.logRequests(s.rateLimit(mux)))
}

func (s *Server) recoverPanics(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		defer func() {
			if recovered := recover(); recovered != nil {
				log.Printf("panic method=%s path=%s remote=%s error=%v stack=%s", r.Method, safeLogPath(r.URL.Path), sanitizeLogValue(r.RemoteAddr), recovered, debug.Stack())
				s3Error(w, http.StatusInternalServerError, "InternalError", "internal server error")
			}
		}()
		next.ServeHTTP(w, r)
	})
}

func (s *Server) metricsHandler(w http.ResponseWriter, r *http.Request) {
	if !s.opsAuthorized(r) {
		writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "unauthorized"})
		return
	}
	w.Header().Set("Content-Type", "text/plain; version=0.0.4")
	fmt.Fprintf(w, "magpie_requests_total %d\n", atomic.LoadUint64(&s.metrics.Requests))
	fmt.Fprintf(w, "magpie_errors_total %d\n", atomic.LoadUint64(&s.metrics.Errors))
	fmt.Fprintf(w, "magpie_response_bytes_total %d\n", atomic.LoadUint64(&s.metrics.ResponseBytes))
}

func (s *Server) opsAuthorized(r *http.Request) bool {
	return s.authorized(r) || s.authorizedS3HeaderPermission(r, "admin")
}

func (cfg Config) withDefaults() Config {
	if len(cfg.S3Keys) == 0 && cfg.S3AccessKeyID != "" && cfg.S3SecretAccessKey != "" {
		cfg.S3Keys = []AccessKey{{ID: cfg.S3AccessKeyID, Secret: cfg.S3SecretAccessKey, Permissions: permissions("read,write,admin")}}
	}
	if cfg.S3Region == "" {
		cfg.S3Region = "auto"
	}
	if cfg.MaxObjectSize <= 0 {
		cfg.MaxObjectSize = defaultMaxObjectSize
	}
	if cfg.PresignedMaxExpiry <= 0 {
		cfg.PresignedMaxExpiry = defaultPresignedMaxExpiry
	}
	if cfg.MultipartMaxAge <= 0 {
		cfg.MultipartMaxAge = defaultMultipartMaxAge
	}
	if cfg.PublicPrefixes == nil {
		cfg.PublicPrefixes = map[string]bool{}
	}
	if cfg.ReplicationMaxJobs <= 0 {
		cfg.ReplicationMaxJobs = defaultReplicationMaxJobs
	}
	if cfg.ReplicationMaxAttempts <= 0 {
		cfg.ReplicationMaxAttempts = defaultReplicationMaxAttempts
	}
	if cfg.MaxRestoreBytes <= 0 {
		cfg.MaxRestoreBytes = defaultMaxRestoreBytes
	}
	return cfg
}

func (s *Server) health(w http.ResponseWriter, r *http.Request) {
	if !s.opsAuthorized(r) {
		writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "unauthorized"})
		return
	}
	disk, err := s.store.DiskStatus()
	if err != nil {
		writeJSON(w, http.StatusServiceUnavailable, map[string]string{"status": "error", "error": err.Error()})
		return
	}
	status := "ok"
	code := http.StatusOK
	if disk.FreeBytes < s.cfg.MinFreeBytes {
		status = "low_disk"
		code = http.StatusServiceUnavailable
	}
	writeJSON(w, code, map[string]any{"status": status, "disk": disk})
}

func (s *Server) object(w http.ResponseWriter, r *http.Request) {
	if !s.authorized(r) {
		writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "unauthorized"})
		return
	}

	key := strings.TrimPrefix(r.URL.Path, objectPrefix)
	if err := ValidateKey(key); err != nil {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": err.Error()})
		return
	}
	bucket, _, _ := strings.Cut(key, "/")
	if !s.bucketAllowed(bucket) {
		writeJSON(w, http.StatusForbidden, map[string]string{"error": "bucket is not allowed"})
		return
	}

	switch r.Method {
	case http.MethodPut:
		s.put(w, r, key)
	case http.MethodGet:
		s.get(w, r, key)
	case http.MethodHead:
		s.head(w, r, key)
	case http.MethodDelete:
		s.delete(w, key)
	default:
		w.Header().Set("Allow", "PUT, GET, HEAD, DELETE")
		writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
	}
}

func (s *Server) put(w http.ResponseWriter, r *http.Request, key string) {
	release, ok := s.acceptWrite(w, r, false)
	if !ok {
		return
	}
	defer release()
	contentType := r.Header.Get("Content-Type")
	info, err := s.store.Put(key, http.MaxBytesReader(w, r.Body, s.cfg.MaxObjectSize), contentType)
	if err != nil {
		if maxBytesError(err) {
			writeJSON(w, http.StatusRequestEntityTooLarge, map[string]string{"error": "object too large"})
			return
		}
		writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "store failed"})
		return
	}
	writeJSON(w, http.StatusCreated, info)
}

func (s *Server) acceptWrite(w http.ResponseWriter, r *http.Request, s3 bool) (func(), bool) {
	size := r.ContentLength
	if r.ContentLength < 0 {
		if s.cfg.MaxObjectSize <= 0 {
			disk, err := s.store.DiskStatus()
			if err != nil || disk.FreeBytes < s.cfg.MinFreeBytes {
				if s3 {
					s3Error(w, http.StatusInsufficientStorage, "InsufficientStorage", "insufficient free disk space for unknown length write")
				} else {
					writeJSON(w, http.StatusInsufficientStorage, map[string]string{"error": "insufficient free disk space for unknown length write"})
				}
				return nil, false
			}
			return func() {}, true
		}
		size = s.cfg.MaxObjectSize
	}
	if s.cfg.MaxObjectSize > 0 && r.ContentLength > s.cfg.MaxObjectSize {
		if s3 {
			s3Error(w, http.StatusRequestEntityTooLarge, "EntityTooLarge", "object too large")
		} else {
			writeJSON(w, http.StatusRequestEntityTooLarge, map[string]string{"error": "object too large"})
		}
		return nil, false
	}
	return s.acceptWriteSize(w, size, s3)
}

func (s *Server) acceptWriteSize(w http.ResponseWriter, size int64, s3 bool) (func(), bool) {
	if size < 0 {
		return func() {}, true
	}
	disk, err := s.store.DiskStatus()
	if err != nil {
		if s3 {
			s3Error(w, http.StatusInternalServerError, "InternalError", "disk status unavailable")
		} else {
			writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "disk status unavailable"})
		}
		return nil, false
	}
	needed := uint64(size)
	s.writeMu.Lock()
	defer s.writeMu.Unlock()
	reserved := s.reservedBytes
	available := uint64(0)
	if disk.FreeBytes > reserved {
		available = disk.FreeBytes - reserved
	}
	if available < s.cfg.MinFreeBytes || needed > available || available-needed < s.cfg.MinFreeBytes {
		if s3 {
			s3Error(w, http.StatusInsufficientStorage, "InsufficientStorage", "insufficient free disk space")
		} else {
			writeJSON(w, http.StatusInsufficientStorage, map[string]string{"error": "insufficient free disk space"})
		}
		return nil, false
	}
	s.reservedBytes += needed
	released := false
	return func() {
		s.writeMu.Lock()
		defer s.writeMu.Unlock()
		if released {
			return
		}
		released = true
		if s.reservedBytes >= needed {
			s.reservedBytes -= needed
		} else {
			s.reservedBytes = 0
		}
	}, true
}

func maxBytesError(err error) bool {
	var maxErr *http.MaxBytesError
	return errors.As(err, &maxErr)
}

func (s *Server) get(w http.ResponseWriter, r *http.Request, key string) {
	obj, err := s.store.Open(key)
	if err != nil {
		status := http.StatusInternalServerError
		if errors.Is(err, ErrNotFound) {
			status = http.StatusNotFound
		}
		writeJSON(w, status, map[string]string{"error": "not found"})
		return
	}
	defer obj.Close()

	contentType := obj.Info.ContentType
	if contentType == "" {
		contentType = "application/octet-stream"
	}
	w.Header().Set("Content-Type", contentType)
	w.Header().Set("X-Object-Key", obj.Info.Key)
	w.Header().Set("X-Object-Size", fmt.Sprintf("%d", obj.Info.Size))
	w.Header().Set("X-Content-Type-Options", "nosniff")
	http.ServeContent(w, r, obj.Info.Key, obj.Info.ModTime, obj)
}

func (s *Server) head(w http.ResponseWriter, r *http.Request, key string) {
	info, err := s.store.Stat(key)
	if err != nil {
		if errors.Is(err, ErrNotFound) {
			w.WriteHeader(http.StatusNotFound)
			return
		}
		w.WriteHeader(http.StatusInternalServerError)
		return
	}

	contentType := info.ContentType
	if contentType == "" {
		contentType = "application/octet-stream"
	}
	w.Header().Set("Content-Type", contentType)
	w.Header().Set("Content-Length", fmt.Sprintf("%d", info.Size))
	w.Header().Set("X-Object-Key", info.Key)
	w.Header().Set("X-Object-Size", fmt.Sprintf("%d", info.Size))
	w.Header().Set("X-Content-Type-Options", "nosniff")
	w.WriteHeader(http.StatusOK)
	_ = r.Body.Close()
}

func (s *Server) delete(w http.ResponseWriter, key string) {
	err := s.store.Delete(key)
	if err != nil && !errors.Is(err, ErrNotFound) {
		writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "delete failed"})
		return
	}
	w.WriteHeader(http.StatusNoContent)
}

func (s *Server) authorized(r *http.Request) bool {
	const prefix = "Bearer "
	value := r.Header.Get("Authorization")
	if !strings.HasPrefix(value, prefix) {
		return false
	}
	provided := strings.TrimPrefix(value, prefix)
	return s.cfg.AuthToken != "" && subtle.ConstantTimeCompare([]byte(provided), []byte(s.cfg.AuthToken)) == 1
}

func (s *Server) bucketAllowed(bucket string) bool {
	return len(s.cfg.AllowedBuckets) == 0 || s.cfg.AllowedBuckets[bucket]
}

func (s *Server) publicReadAllowed(key string) bool {
	for prefix := range s.cfg.PublicPrefixes {
		prefix = strings.TrimSuffix(prefix, "/")
		if !strings.Contains(prefix, "/") {
			continue
		}
		if key == prefix || strings.HasPrefix(key, prefix+"/") {
			return true
		}
	}
	return false
}

func (s *Server) replicationAllowed(r *http.Request) bool {
	if r.Header.Get("X-Magpie-Replication") != "1" {
		return true
	}
	if s.cfg.ReplicationSecret == "" {
		return false
	}
	return subtle.ConstantTimeCompare([]byte(r.Header.Get("X-Magpie-Replication-Secret")), []byte(s.cfg.ReplicationSecret)) == 1
}

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

type statusWriter struct {
	http.ResponseWriter
	status int
	bytes  int
}

func (w *statusWriter) WriteHeader(status int) {
	w.status = status
	w.ResponseWriter.WriteHeader(status)
}

func (w *statusWriter) Write(b []byte) (int, error) {
	if w.status == 0 {
		w.status = http.StatusOK
	}
	n, err := w.ResponseWriter.Write(b)
	w.bytes += n
	return n, err
}

func (s *Server) logRequests(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		recorder := &statusWriter{ResponseWriter: w}
		next.ServeHTTP(recorder, r)
		status := recorder.status
		if status == 0 {
			status = http.StatusOK
		}
		atomic.AddUint64(&s.metrics.Requests, 1)
		if recorder.bytes > 0 {
			atomic.AddUint64(&s.metrics.ResponseBytes, uint64(recorder.bytes))
		}
		if status >= 400 {
			atomic.AddUint64(&s.metrics.Errors, 1)
		}
		log.Printf(
			"method=%s path=%s status=%d bytes=%d duration_ms=%d remote=%s",
			r.Method,
			safeLogPath(r.URL.Path),
			status,
			recorder.bytes,
			time.Since(start).Milliseconds(),
			sanitizeLogValue(r.RemoteAddr),
		)
	})
}

func (s *Server) rateLimit(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !s.rateLimiter.Allow(s.rateLimitKey(r)) {
			writeJSON(w, http.StatusTooManyRequests, map[string]string{"error": "rate limited"})
			return
		}
		next.ServeHTTP(w, r)
	})
}

func (s *Server) rateLimitKey(r *http.Request) string {
	if s.cfg.TrustProxyHeaders {
		if forwarded := r.Header.Get("X-Forwarded-For"); forwarded != "" {
			candidate := strings.TrimSpace(strings.Split(forwarded, ",")[0])
			if net.ParseIP(candidate) != nil {
				return candidate
			}
		}
		if realIP := strings.TrimSpace(r.Header.Get("X-Real-IP")); net.ParseIP(realIP) != nil {
			return realIP
		}
	}
	return r.RemoteAddr
}

type Metrics struct {
	Requests      uint64
	Errors        uint64
	ResponseBytes uint64
}

type RateLimiter struct {
	limit   int
	mu      sync.Mutex
	buckets map[string]rateBucket
}

type rateBucket struct {
	minute int64
	count  int
}

func NewRateLimiter(limit int) *RateLimiter {
	return &RateLimiter{limit: limit, buckets: map[string]rateBucket{}}
}

func (r *RateLimiter) Allow(remote string) bool {
	if r.limit <= 0 {
		return true
	}
	ip, _, err := net.SplitHostPort(remote)
	if err != nil {
		ip = remote
	}
	minute := time.Now().Unix() / 60
	r.mu.Lock()
	defer r.mu.Unlock()
	for candidate, bucket := range r.buckets {
		if bucket.minute < minute-1 {
			delete(r.buckets, candidate)
		}
	}
	bucket := r.buckets[ip]
	if bucket.minute != minute {
		bucket = rateBucket{minute: minute}
	}
	bucket.count++
	r.buckets[ip] = bucket
	return bucket.count <= r.limit
}

func safeLogPath(path string) string {
	if len(path) <= 256 {
		return sanitizeLogValue(path)
	}
	return sanitizeLogValue(path[:256]) + "..."
}

func sanitizeLogValue(value string) string {
	return strings.Map(func(r rune) rune {
		if r < 0x20 || r == 0x7f {
			return ' '
		}
		return r
	}, value)
}

type ObjectInfo struct {
	Key         string    `json:"key"`
	Size        int64     `json:"size"`
	ContentType string    `json:"content_type"`
	ModTime     time.Time `json:"mod_time"`
	ETag        string    `json:"etag"`
}

type DiskStatus struct {
	TotalBytes uint64 `json:"total_bytes"`
	FreeBytes  uint64 `json:"free_bytes"`
	UsedBytes  uint64 `json:"used_bytes"`
}

type StoreStats struct {
	ObjectCount      int64                 `json:"object_count"`
	TotalBytes       int64                 `json:"total_bytes"`
	Disk             DiskStatus            `json:"disk"`
	ReplicationQueue ReplicationQueueStats `json:"replication_queue"`
}

type ScrubReport struct {
	CheckedObjects int      `json:"checked_objects"`
	MissingKeys    []string `json:"missing_keys"`
	CorruptKeys    []string `json:"corrupt_keys"`
	OrphanKeys     []string `json:"orphan_keys"`
}

type ListPage struct {
	Objects   []ObjectInfo
	NextAfter string
}