package server

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	"git.n1ko.dev/Niko/niko_trust/pkg/identity"
	"git.n1ko.dev/Niko/niko_trust/pkg/protocol"
	"git.n1ko.dev/Niko/niko_trust/pkg/tce"
	"git.n1ko.dev/Niko/niko_trust/pkg/transport"
)

// Server is the trust relay: it stores signed objects and brokers
// authentication, without ever holding a signing key or making authorization
// decisions (INV-1, INV-5).
type Server struct {
	store    *Store
	audience string
	metrics  *Metrics
	ready    bool

	maxBodyBytes int64
	challengeTTL time.Duration
	sessionTTL   time.Duration
	trustProxy   bool

	putLimiter       *ipLimiter
	challengeLimiter *ipLimiter

	mu         sync.Mutex
	challenges map[string]time.Time // challenge hex -> expiry
	sessions   map[string]session   // session token -> session

	// ws streams newly stored envelopes to subscribed sessions.
	ws *wsHubState
}

type session struct {
	identity string
	scope    string
	expiry   time.Time
}

// New builds a relay from the given Config. Any zero-valued field falls back to
// DefaultConfig, so callers may pass a partially-populated Config.
func New(cfg Config) *Server {
	if cfg.ListenAddr == "" {
		cfg.ListenAddr = ":8080"
	}
	if cfg.Audience == "" {
		cfg.Audience = "trust.n1ko.dev"
	}
	if cfg.LogLevel == "" {
		cfg.LogLevel = "info"
	}
	if cfg.MaxBodyBytes <= 0 {
		cfg.MaxBodyBytes = tce.MaxClaimTCE*2 + 1024
	}
	if cfg.PutLimit <= 0 {
		cfg.PutLimit = 60
	}
	if cfg.PutWindow <= 0 {
		cfg.PutWindow = time.Minute
	}
	if cfg.ChallengeLimit <= 0 {
		cfg.ChallengeLimit = 30
	}
	if cfg.ChallengeWindow <= 0 {
		cfg.ChallengeWindow = time.Minute
	}
	if cfg.MaxPerSubject <= 0 {
		cfg.MaxPerSubject = 1000
	}
	if cfg.ChallengeTTL <= 0 {
		cfg.ChallengeTTL = 5 * time.Minute
	}
	if cfg.SessionTTL <= 0 {
		cfg.SessionTTL = 30 * time.Minute
	}
	s := &Server{
		store:            NewStore(cfg.DataDir, cfg.MaxPerSubject),
		audience:         cfg.Audience,
		metrics:          NewMetrics(),
		maxBodyBytes:     cfg.MaxBodyBytes,
		challengeTTL:     cfg.ChallengeTTL,
		sessionTTL:       cfg.SessionTTL,
		trustProxy:       cfg.TrustProxy,
		challenges:       make(map[string]time.Time),
		sessions:         make(map[string]session),
		ws:               newWSHub(),
		putLimiter:       newIPLimiter(cfg.PutLimit, cfg.PutWindow),
		challengeLimiter: newIPLimiter(cfg.ChallengeLimit, cfg.ChallengeWindow),
	}
	s.ready = true
	return s
}

// Store exposes the underlying object store.
func (s *Server) Store() *Store { return s.store }

// Handler returns the HTTP handler implementing the JSON transport.
func (s *Server) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("POST /v1/objects", s.rateLimitPut)
	mux.HandleFunc("GET /v1/objects/{id}", s.handleGet)
	mux.HandleFunc("GET /v1/objects", s.handleObjectsBatch)
	mux.HandleFunc("GET /v1/claims", s.handleClaims)
	mux.HandleFunc("GET /v1/requests", s.handleRequests)
	mux.HandleFunc("GET /v1/responses", s.handleResponses)
	mux.HandleFunc("GET /v1/revocations", s.handleRevocations)
	mux.HandleFunc("GET /v1/config", s.handleConfig)
	mux.HandleFunc("GET /v1/metrics", s.handleMetrics)
	mux.HandleFunc("GET /v1/healthz", s.handleHealth)
	mux.HandleFunc("GET /v1/readyz", s.handleReady)
	mux.HandleFunc("POST /v1/auth/challenge", s.rateLimitChallenge)
	mux.HandleFunc("POST /v1/auth/assert", s.handleAssert)
	mux.HandleFunc("GET /v1/ws", s.handleWS)
	return mux
}

func (s *Server) rateLimitPut(w http.ResponseWriter, r *http.Request) {
	if !s.putLimiter.allow(clientIP(r, s.trustProxy)) {
		writeErr(w, http.StatusTooManyRequests, "rate limited")
		return
	}
	s.HandlePut(w, r)
}

func (s *Server) rateLimitChallenge(w http.ResponseWriter, r *http.Request) {
	if !s.challengeLimiter.allow(clientIP(r, s.trustProxy)) {
		writeErr(w, http.StatusTooManyRequests, "rate limited")
		return
	}
	s.handleChallenge(w, r)
}

func (s *Server) handleMetrics(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/plain; version=0.0.4")
	s.metrics.Write(w)
}

func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	_ = json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
}

func (s *Server) handleReady(w http.ResponseWriter, r *http.Request) {
	if s.ready {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_ = json.NewEncoder(w).Encode(map[string]string{"status": "ready"})
		return
	}
	writeErr(w, http.StatusServiceUnavailable, "not ready")
}

func (s *Server) HandlePut(w http.ResponseWriter, r *http.Request) {
	req, ok := s.decodeWire(w, r)
	if !ok {
		return
	}
	if len(req.TCE) == 0 {
		writeErr(w, http.StatusBadRequest, "missing tce")
		return
	}
	objectID, created, err := s.store.Put(req.TCE, req.Signature, req.ObjectID)
	if err != nil {
		s.metrics.inc(&s.metrics.objectsRejected)
		writeErr(w, http.StatusUnprocessableEntity, err.Error())
		return
	}
	s.metrics.inc(&s.metrics.objectsStored)
	if created {
		if typ, obj, derr := transport.DecodeObject(req.TCE); derr == nil {
			s.metrics.incStoredType(typ)
			s.hubBroadcastObj(typ, obj, req.TCE, req.Signature, objectID)
		}
	}
	writeJSON(w, http.StatusOK, map[string]string{"object_id": objectID})
}

func (s *Server) handleGet(w http.ResponseWriter, r *http.Request) {
	id := r.PathValue("id")
	env, err := s.store.Get(id)
	if err != nil {
		writeErr(w, http.StatusNotFound, err.Error())
		return
	}
	writeJSON(w, http.StatusOK, env)
}

// scopeOK reports whether a session may perform an action requiring the given
// scope. Besides the exact scope, a session granted the generic "read" or "*"
// capability may read any read endpoint.
func scopeOK(ses session, required string) bool {
	return ses.scope == required || ses.scope == "read" || ses.scope == "*"
}

// paginate applies optional limit/offset/after query params to a result
// list. A non-positive limit means "no cap". `after` keeps only entries
// whose content id is lexicographically greater, giving clients a stable
// cursor across concurrent inserts.
func paginate(in []*transport.Envelope, limit, offset int, after string) []*transport.Envelope {
	if after != "" {
		kept := in[:0]
		for _, env := range in {
			if tce.ComputeID(env.TCE).String() > after {
				kept = append(kept, env)
			}
		}
		in = kept
	}
	if offset > 0 {
		if offset >= len(in) {
			return nil
		}
		in = in[offset:]
	}
	if limit > 0 && limit < len(in) {
		in = in[:limit]
	}
	return in
}

// listParams reads limit/offset/after from the query string.
func listParams(r *http.Request) (limit, offset int, after string) {
	if v := r.URL.Query().Get("limit"); v != "" {
		if n, err := strconv.Atoi(v); err == nil {
			limit = n
		}
	}
	if v := r.URL.Query().Get("offset"); v != "" {
		if n, err := strconv.Atoi(v); err == nil {
			offset = n
		}
	}
	return limit, offset, r.URL.Query().Get("after")
}

const maxBatchIDs = 100

// handleObjectsBatch serves GET /v1/objects?ids=a,b,c — one round trip for
// consumers that already know which objects they need. Missing ids are
// omitted rather than erroring.
func (s *Server) handleObjectsBatch(w http.ResponseWriter, r *http.Request) {
	raw := r.URL.Query().Get("ids")
	if raw == "" {
		writeErr(w, http.StatusBadRequest, "missing ids")
		return
	}
	parts := strings.Split(raw, ",")
	if len(parts) > maxBatchIDs {
		writeErr(w, http.StatusBadRequest, "too many ids")
		return
	}
	clean := make([]string, 0, len(parts))
	for _, id := range parts {
		id = strings.TrimSpace(id)
		if id == "" {
			continue
		}
		if _, err := tce.ParseID(id); err != nil {
			writeErr(w, http.StatusBadRequest, "malformed object id: "+id)
			return
		}
		clean = append(clean, id)
	}
	found, err := s.store.GetMany(clean)
	if err != nil {
		writeErr(w, http.StatusInternalServerError, err.Error())
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{"objects": found})
}

func (s *Server) handleClaims(w http.ResponseWriter, r *http.Request) {
	ses, ok := s.sessionByToken(r)
	if !ok {
		writeErr(w, http.StatusUnauthorized, "authentication required")
		return
	}
	if !scopeOK(ses, "read:claims") {
		writeErr(w, http.StatusForbidden, "insufficient scope")
		return
	}
	subject := r.URL.Query().Get("subject")
	if subject == "" {
		writeErr(w, http.StatusBadRequest, "missing subject")
		return
	}
	list, err := s.store.ClaimsBySubject(subject)
	if err != nil {
		writeErr(w, http.StatusInternalServerError, err.Error())
		return
	}
	lim, off, after := listParams(r)
	writeJSON(w, http.StatusOK, map[string]any{"claims": paginate(list, lim, off, after)})
}

func (s *Server) handleRequests(w http.ResponseWriter, r *http.Request) {
	ses, ok := s.sessionByToken(r)
	if !ok {
		writeErr(w, http.StatusUnauthorized, "authentication required")
		return
	}
	if !scopeOK(ses, "read:requests") {
		writeErr(w, http.StatusForbidden, "insufficient scope")
		return
	}
	recipient := r.URL.Query().Get("recipient")
	if recipient == "" {
		writeErr(w, http.StatusBadRequest, "missing recipient")
		return
	}
	list, err := s.store.RequestsByRecipient(recipient)
	if err != nil {
		writeErr(w, http.StatusInternalServerError, err.Error())
		return
	}
	lim, off, after := listParams(r)
	writeJSON(w, http.StatusOK, map[string]any{"requests": paginate(list, lim, off, after)})
}

func (s *Server) handleResponses(w http.ResponseWriter, r *http.Request) {
	ses, ok := s.sessionByToken(r)
	if !ok {
		writeErr(w, http.StatusUnauthorized, "authentication required")
		return
	}
	if !scopeOK(ses, "read:responses") {
		writeErr(w, http.StatusForbidden, "insufficient scope")
		return
	}
	req := r.URL.Query().Get("request")
	if req == "" {
		writeErr(w, http.StatusBadRequest, "missing request")
		return
	}
	list, err := s.store.ResponsesForRequest(req)
	if err != nil {
		writeErr(w, http.StatusInternalServerError, err.Error())
		return
	}
	lim, off, after := listParams(r)
	writeJSON(w, http.StatusOK, map[string]any{"responses": paginate(list, lim, off, after)})
}

func (s *Server) handleRevocations(w http.ResponseWriter, r *http.Request) {
	ses, ok := s.sessionByToken(r)
	if !ok {
		writeErr(w, http.StatusUnauthorized, "authentication required")
		return
	}
	if !scopeOK(ses, "read:revocations") {
		writeErr(w, http.StatusForbidden, "insufficient scope")
		return
	}
	claim := r.URL.Query().Get("claim")
	if claim == "" {
		writeErr(w, http.StatusBadRequest, "missing claim")
		return
	}
	list, err := s.store.RevocationsForClaim(claim)
	if err != nil {
		writeErr(w, http.StatusInternalServerError, err.Error())
		return
	}
	lim, off, after := listParams(r)
	writeJSON(w, http.StatusOK, map[string]any{"revocations": paginate(list, lim, off, after)})
}

func (s *Server) handleConfig(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, http.StatusOK, map[string]string{"audience": s.audience})
}

func (s *Server) handleChallenge(w http.ResponseWriter, r *http.Request) {
	ch := make([]byte, tce.ChallengeSize)
	if _, err := rand.Read(ch); err != nil {
		writeErr(w, http.StatusInternalServerError, "challenge")
		return
	}
	chHex := hex.EncodeToString(ch)
	s.mu.Lock()
	// Drop any challenges that have lapsed without being used.
	nowTs := now()
	for k, exp := range s.challenges {
		if exp.Before(nowTs) {
			delete(s.challenges, k)
		}
	}
	s.challenges[chHex] = nowTs.Add(s.challengeTTL)
	s.mu.Unlock()
	s.metrics.inc(&s.metrics.challengesIssued)
	writeJSON(w, http.StatusOK, map[string]string{"challenge": chHex})
}

func (s *Server) handleAssert(w http.ResponseWriter, r *http.Request) {
	env, ok := s.decodeEnvelope(w, r)
	if !ok {
		return
	}
	// Strict-decode and verify the assertion, bound to this server's audience.
	if _, _, err := transport.DecodeObject(env.TCE); err != nil {
		writeErr(w, http.StatusUnprocessableEntity, "rejected: "+err.Error())
		return
	}
	assert, err := protocol.DecodeAuthAssertion(env.TCE)
	if err != nil {
		writeErr(w, http.StatusUnprocessableEntity, "rejected")
		return
	}
	if _, err := protocol.VerifyAuthAssertion(env.TCE, env.Signature, s.audience); err != nil {
		writeErr(w, http.StatusUnauthorized, "auth failed: "+err.Error())
		return
	}
	// The assertion must be bound to a challenge this server actually issued,
	// and each challenge is single-use: consuming it here blocks replay of a
	// captured assertion.
	chHex := hex.EncodeToString(assert.Challenge)
	s.mu.Lock()
	expiry, issued := s.challenges[chHex]
	if issued {
		delete(s.challenges, chHex)
	}
	s.mu.Unlock()
	if !issued || expiry.Before(now()) {
		s.metrics.inc(&s.metrics.assertionsFailed)
		writeErr(w, http.StatusUnauthorized, "unknown or expired challenge")
		return
	}
	// The assertion is valid for this audience; mint an opaque session token.
	// The server holds no signing key (INV-1): the token is a random reference
	// to the verified identity, stored server-side.
	tok := make([]byte, 32)
	if _, err := rand.Read(tok); err != nil {
		writeErr(w, http.StatusInternalServerError, "session")
		return
	}
	tokHex := hex.EncodeToString(tok)
	id := identityFromPubKey(assert.PubKey)
	s.mu.Lock()
	s.sessions[tokHex] = session{identity: id, scope: assert.Scope, expiry: now().Add(s.sessionTTL)}
	s.metrics.sessionsActive.Add(1)
	s.mu.Unlock()
	s.metrics.inc(&s.metrics.assertionsOK)
	writeJSON(w, http.StatusOK, map[string]string{
		"session_token": tokHex,
		"identity":      id,
		"scope":         assert.Scope,
	})
}

// sessionByToken returns a valid, unexpired session for the request's bearer
// token, if any. Expired sessions are purged on access.
func (s *Server) sessionByToken(r *http.Request) (session, bool) {
	tok := bearerToken(r)
	if tok == "" {
		return session{}, false
	}
	s.mu.Lock()
	ses, ok := s.sessions[tok]
	if ok {
		if ses.expiry.Before(now()) {
			delete(s.sessions, tok)
			s.metrics.sessionsActive.Add(-1)
			ok = false
		}
	}
	s.mu.Unlock()
	return ses, ok
}

// bearerToken extracts a session token from the Authorization header or the
// ?token= query parameter.
func bearerToken(r *http.Request) string {
	if h := r.Header.Get("Authorization"); strings.HasPrefix(h, "Bearer ") {
		return strings.TrimSpace(h[len("Bearer "):])
	}
	return r.URL.Query().Get("token")
}

func identityFromPubKey(pub []byte) string {
	id, err := identity.FromPubKey(pub)
	if err != nil {
		return ""
	}
	return id.Address().String()
}

// wireRequest is the JSON body of endpoints that carry a signed envelope:
// object storage and authentication. The optional pow field carries a solved
// admission challenge; it is transport metadata and never reaches TCE bytes.
type wireRequest struct {
	TCE       []byte `json:"tce"`
	Signature []byte `json:"signature"`
	ObjectID  string `json:"object_id,omitempty"`
}

// decodeWire reads a JSON request body from the request, enforcing the
// server-wide body cap first.
func (s *Server) decodeWire(w http.ResponseWriter, r *http.Request) (*wireRequest, bool) {
	r.Body = http.MaxBytesReader(w, r.Body, s.maxBodyBytes)
	var req wireRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		var maxErr *http.MaxBytesError
		if errors.As(err, &maxErr) {
			writeErr(w, http.StatusRequestEntityTooLarge, "payload too large")
		} else {
			writeErr(w, http.StatusBadRequest, "bad envelope")
		}
		return nil, false
	}
	return &req, true
}

// decodeEnvelope adapts decodeWire to the plain envelope form for callers
// that do not care about admission control.
func (s *Server) decodeEnvelope(w http.ResponseWriter, r *http.Request) (*transport.Envelope, bool) {
	req, ok := s.decodeWire(w, r)
	if !ok {
		return nil, false
	}
	return &transport.Envelope{TCE: req.TCE, Signature: req.Signature}, true
}

// readCapped reads the request body under the server-wide cap.
func (s *Server) readCapped(w http.ResponseWriter, r *http.Request) []byte {
	raw, err := io.ReadAll(http.MaxBytesReader(w, r.Body, s.maxBodyBytes))
	if err != nil {
		return nil
	}
	return raw
}

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

func writeErr(w http.ResponseWriter, code int, msg string) {
	writeJSON(w, code, map[string]string{"error": msg})
}