// Copyright (C) 2026 Dcntrlized
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

package owner

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"time"

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

const (
	// AccessTokenTTL is the stateless access-token lifetime.
	AccessTokenTTL = 15 * time.Minute
	// RefreshTokenTTL is the persisted refresh-token lifetime.
	RefreshTokenTTL = 7 * 24 * time.Hour
	// Issuer tags every minted token; verification rejects others.
	Issuer = "maestro"

	claimType = "type"
	claimKid  = "kid"
	claimJti  = "jti"

	typeAccess  = "access"
	typeRefresh = "refresh"
)

var (
	ErrInvalidToken = errors.New("invalid token")
	ErrWrongType    = errors.New("token has the wrong type")
)

// TokenPair is one login/rotation result: the cookie payloads plus the refresh
// jti (its raw value lives only in the cookies and the caller's response).
type TokenPair struct {
	Access     string
	Refresh    string
	RefreshJTI string
	AccessTTL  time.Duration
	RefreshTTL time.Duration
	Username   string
	Role       string
}

// TokenService mints and verifies the HS256 access/refresh session pair. The
// signing secret is a per-instance random key on disk (`session.key`, chmod
// 600); kid = sha256(secret)[:16] so a rotated secret keeps old access tokens
// valid through their remaining TTL (grace window).
type TokenService struct {
	dir string

	key []byte
	kid string
}

// NewTokenService loads or creates dir/session.key.
func NewTokenService(dir string) (*TokenService, error) {
	if err := os.MkdirAll(dir, 0o700); err != nil {
		return nil, fmt.Errorf("create token directory: %w", err)
	}
	path := filepath.Join(dir, "session.key")
	key, err := os.ReadFile(path)
	if errors.Is(err, os.ErrNotExist) {
		key = make([]byte, 32)
		if _, err := rand.Read(key); err != nil {
			return nil, fmt.Errorf("generate session key: %w", err)
		}
		if err := os.WriteFile(path, key, 0o600); err != nil {
			return nil, fmt.Errorf("persist session key: %w", err)
		}
	} else if err != nil {
		return nil, fmt.Errorf("read session key: %w", err)
	} else if len(key) < 32 {
		return nil, fmt.Errorf("session key at %s is too short", path)
	}

	sum := sha256.Sum256(key)
	return &TokenService{dir: dir, key: key, kid: hex.EncodeToString(sum[:])[:16]}, nil
}

// Kid returns the active signing-key identifier.
func (t *TokenService) Kid() string { return t.kid }

func (t *TokenService) claims(sub, typ, role string, ttl time.Duration, jti string) jwt.MapClaims {
	now := time.Now()
	c := jwt.MapClaims{
		"sub":     sub,
		"iss":     Issuer,
		"iat":     now.Unix(),
		"exp":     now.Add(ttl).Unix(),
		claimType: typ,
		claimKid:  t.kid,
	}
	if role != "" {
		c["role"] = role
	}
	if jti != "" {
		c[claimJti] = jti
	}
	return c
}

func (t *TokenService) sign(c jwt.MapClaims) (string, error) {
	return jwt.NewWithClaims(jwt.SigningMethodHS256, c).SignedString(t.key)
}

// MintPair issues a fresh access/refresh pair for the owner. The refresh jti
// must be recorded by the session store before the cookie is served.
func (t *TokenService) MintPair(username, role string) (*TokenPair, error) {
	jti := make([]byte, 16)
	if _, err := rand.Read(jti); err != nil {
		return nil, fmt.Errorf("generate jti: %w", err)
	}
	raw := hex.EncodeToString(jti)

	access, err := t.sign(t.claims(username, typeAccess, role, AccessTokenTTL, ""))
	if err != nil {
		return nil, fmt.Errorf("sign access token: %w", err)
	}
	refresh, err := t.sign(t.claims(username, typeRefresh, "", RefreshTokenTTL, raw))
	if err != nil {
		return nil, fmt.Errorf("sign refresh token: %w", err)
	}
	return &TokenPair{
		Access: access, Refresh: refresh, RefreshJTI: raw,
		AccessTTL: AccessTokenTTL, RefreshTTL: RefreshTokenTTL,
		Username: username, Role: role,
	}, nil
}

// parse validates signature (constant-time HMAC compare inside the library),
// issuer and expiry, returning the claims.
func (t *TokenService) parse(token string) (jwt.MapClaims, error) {
	parsed, err := jwt.Parse(token, func(tk *jwt.Token) (interface{}, error) {
		if _, ok := tk.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, ErrInvalidToken
		}
		return t.key, nil
	}, jwt.WithValidMethods([]string{jwt.SigningMethodHS256.Alg()}), jwt.WithIssuer(Issuer), jwt.WithExpirationRequired())
	if err != nil || !parsed.Valid {
		return nil, ErrInvalidToken
	}
	claims, ok := parsed.Claims.(jwt.MapClaims)
	if !ok {
		return nil, ErrInvalidToken
	}
	if kid, _ := claims[claimKid].(string); kid != "" && kid != t.kid && !hmac.Equal([]byte(kid), []byte(t.kid)) {
		// A different kid means an unknown secret generation — refuse. The
		// same-kid path is the normal case; empty kid predates kid tagging.
		return nil, ErrInvalidToken
	}
	return claims, nil
}

// VerifyAccess validates an access token and returns the username + role.
func (t *TokenService) VerifyAccess(token string) (username, role string, err error) {
	claims, err := t.parse(token)
	if err != nil {
		return "", "", err
	}
	if typ, _ := claims[claimType].(string); typ != typeAccess {
		return "", "", ErrWrongType
	}
	sub, _ := claims["sub"].(string)
	role, _ = claims["role"].(string)
	if sub == "" {
		return "", "", ErrInvalidToken
	}
	return sub, role, nil
}

// VerifyRefresh validates a refresh token and returns username + jti.
func (t *TokenService) VerifyRefresh(token string) (username, jti string, err error) {
	claims, err := t.parse(token)
	if err != nil {
		return "", "", err
	}
	if typ, _ := claims[claimType].(string); typ != typeRefresh {
		return "", "", ErrWrongType
	}
	sub, _ := claims["sub"].(string)
	jti, _ = claims[claimJti].(string)
	if sub == "" || jti == "" {
		return "", "", ErrInvalidToken
	}
	return sub, jti, nil
}