package gitrpc

import (
	"bytes"
	"context"
	"crypto/ed25519"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"simplegit/common"
	"strconv"
	"sync"
	"time"

	"github.com/golang-jwt/jwt/v5"
)

type Claims struct {
	jwt.RegisteredClaims
	Username string `json:"username"`
}

func (c *Claims) UserID() int64 {
	id, _ := strconv.ParseInt(c.Subject, 10, 64)
	return id
}

type ctxKey struct{}

func WithClaims(ctx context.Context, c *Claims) context.Context {
	return context.WithValue(ctx, ctxKey{}, c)
}

func FromContext(ctx context.Context) *Claims {
	c, _ := ctx.Value(ctxKey{}).(*Claims)
	return c
}

type JWK struct {
	Kty string `json:"kty"`
	Crv string `json:"crv"`
	Alg string `json:"alg"`
	Kid string `json:"kid"`
	X   string `json:"x"`
}

type JWKS struct {
	Keys []JWK `json:"keys"`
}

func publicKeyFromJWK(k JWK) (ed25519.PublicKey, error) {
	if k.Kty != "OKP" || k.Crv != "Ed25519" {
		return nil, fmt.Errorf("unsupported jwk: kty=%q crv=%q", k.Kty, k.Crv)
	}
	x, err := base64.RawURLEncoding.DecodeString(k.X)
	if err != nil {
		return nil, fmt.Errorf("decode jwk x: %w", err)
	}
	if len(x) != ed25519.PublicKeySize {
		return nil, fmt.Errorf("bad ed25519 key length: %d", len(x))
	}
	return ed25519.PublicKey(x), nil
}

type jwksSource struct {
	authURL string
	client  *http.Client

	// self-registration: when non-empty, every jwks refresh also announces this
	// service to the console (RBAC center) via X-Service-* headers, so the
	// console can discover and dial this gitrpc server. The console is the
	// service-discovery hub; the only side channel it needs is the jwks fetch
	// every consumer already makes.
	selfName string
	selfKind string
	selfAddr string

	mu   sync.RWMutex
	keys map[string]ed25519.PublicKey
}

func newJWKSSource(authURL, selfName, selfKind, selfAddr string) *jwksSource {
	return &jwksSource{
		authURL:  authURL,
		client:   &http.Client{Timeout: 10 * time.Second},
		selfName: selfName,
		selfKind: selfKind,
		selfAddr: selfAddr,
		keys:     map[string]ed25519.PublicKey{},
	}
}

func (s *jwksSource) refresh() error {
	req, err := http.NewRequest(http.MethodGet, s.authURL+"/.well-known/jwks.json", nil)
	if err != nil {
		return fmt.Errorf("build jwks request: %w", err)
	}
	// Announce ourselves so the console registers this service's address. The
	// console answers jwks AND records the caller as a reachable service.
	if s.selfName != "" && s.selfAddr != "" {
		req.Header.Set("X-Service-Name", s.selfName)
		req.Header.Set("X-Service-Kind", s.selfKind)
		req.Header.Set("X-Service-Addr", s.selfAddr)
	}
	resp, err := s.client.Do(req)
	if err != nil {
		return fmt.Errorf("fetch jwks: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("jwks fetch status %d", resp.StatusCode)
	}
	var set JWKS
	if err := json.NewDecoder(resp.Body).Decode(&set); err != nil {
		return fmt.Errorf("decode jwks: %w", err)
	}
	keys := make(map[string]ed25519.PublicKey, len(set.Keys))
	for _, k := range set.Keys {
		pub, err := publicKeyFromJWK(k)
		if err != nil {
			continue
		}
		keys[k.Kid] = pub
	}
	if len(keys) == 0 {
		return errors.New("jwks has no usable ed25519 keys")
	}
	s.mu.Lock()
	s.keys = keys
	s.mu.Unlock()
	return nil
}

func (s *jwksSource) resolve(kid string) (ed25519.PublicKey, error) {
	s.mu.RLock()
	pub, ok := s.keys[kid]
	s.mu.RUnlock()
	if ok {
		return pub, nil
	}

	if err := s.refresh(); err != nil {
		return nil, fmt.Errorf("refetch jwks for kid %q: %w", kid, err)
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	pub, ok = s.keys[kid]
	if !ok {
		return nil, fmt.Errorf("kid %q not in jwks", kid)
	}
	return pub, nil
}

type Verifier struct {
	resolve func(kid string) (ed25519.PublicKey, error)
}

func NewJWKSVerifier(authURL, selfName, selfKind, selfAddr string) (*Verifier, error) {
	src := newJWKSSource(authURL, selfName, selfKind, selfAddr)
	if err := src.refresh(); err != nil {
		return nil, err
	}
	return &Verifier{resolve: src.resolve}, nil
}

func (v *Verifier) Parse(tokenStr string) (*Claims, error) {
	var c Claims
	_, err := jwt.ParseWithClaims(tokenStr, &c, func(t *jwt.Token) (any, error) {
		if t.Method != jwt.SigningMethodEdDSA {
			return nil, fmt.Errorf("unexpected signing method: %v", t.Header["alg"])
		}
		kid, _ := t.Header["kid"].(string)
		pub, err := v.resolve(kid)
		if err != nil {
			return nil, err
		}
		return pub, nil
	}, jwt.WithExpirationRequired(), jwt.WithIssuer("simpleconsole"))
	if err != nil {
		return nil, err
	}
	return &c, nil
}

type MiddlewareClient struct {
	baseURL  string
	http     *http.Client
	verifier *Verifier
}

// NewMiddleware builds the console RBAC client. selfName/selfKind/selfAddr, when
// non-empty, are announced to the console on every jwks refresh so the console
// registers this server as a reachable service (service discovery). Pass the
// gitrpc HTTP address the console should dial back.
func NewMiddleware(baseURL, selfName, selfKind, selfAddr string) *MiddlewareClient {
	verifier, err := NewJWKSVerifier(baseURL, selfName, selfKind, selfAddr)
	if err != nil {
		return nil
	}
	return &MiddlewareClient{
		baseURL:  baseURL,
		http:     &http.Client{Timeout: 4 * time.Second},
		verifier: verifier,
	}
}

// ParseToken verifies a console-issued JWT and returns its claims. An empty
// token returns (nil, nil) so the caller can distinguish "no credential"
// (anonymous) from "credential present but invalid" (non-nil error): the HTTP
// and SSH access planes use this to decide between anonymous fallthrough and a
// 401. skip-auth callers never invoke this (middleware is nil).
func (c *MiddlewareClient) ParseToken(tokenStr string) (*Claims, error) {
	if tokenStr == "" {
		return nil, nil
	}
	return c.verifier.Parse(tokenStr)
}

func (c *MiddlewareClient) Authz(ctx context.Context, userID int64, repo string, perm common.Perm) error {

	body, _ := json.Marshal(struct {
		Repo   string `json:"repo"`
		UserID string `json:"userId"`
		Perm   string `json:"perm"`
	}{repo, strconv.FormatInt(userID, 10), string(perm)})

	resp, err := c.postConnect(ctx, "/simpleconsole.v1.InternalService/Authz", body)
	if err != nil {
		return fmt.Errorf("authz call: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("authz status %d for %s %s", resp.StatusCode, repo, perm)
	}
	var out struct {
		Allow bool `json:"allow"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return fmt.Errorf("decode authz: %w", err)
	}
	if !out.Allow {
		return ErrDenied
	}
	return nil
}

var ErrDenied = errors.New("access denied")

// ErrInvalidToken marks a credential that was present but could not be verified
// (expired, unsigned, malformed). It is distinct from anonymous (no credential
// at all) and from ErrDenied (verified identity lacking permission).
var ErrInvalidToken = errors.New("invalid or expired token")

func (c *MiddlewareClient) Authn(ctx context.Context, fingerprint string) (userID int64, username string, err error) {

	body, _ := json.Marshal(struct {
		Fingerprint string `json:"fingerprint"`
	}{fingerprint})

	resp, err := c.postConnect(ctx, "/simpleconsole.v1.InternalService/Authn", body)
	if err != nil {
		return 0, "", fmt.Errorf("authn call: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return 0, "", fmt.Errorf("authn status %d", resp.StatusCode)
	}

	var out struct {
		UserID   string `json:"userId"`
		Username string `json:"username"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		return 0, "", fmt.Errorf("decode authn: %w", err)
	}
	uid, err := strconv.ParseInt(out.UserID, 10, 64)
	if err != nil {
		return 0, "", fmt.Errorf("decode authn userId: %w", err)
	}
	return uid, out.Username, nil
}

func (c *MiddlewareClient) postConnect(ctx context.Context, procedure string, body []byte) (*http.Response, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+procedure, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Connect-Protocol-Version", "1")
	return c.http.Do(req)
}