package capabilities

import (
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"
)

const (
	issuer   = "agentteam"
	audience = "agentteam-internal-app-api"
)

var (
	ErrInvalidToken = errors.New("invalid app capability token")
	ErrExpiredToken = errors.New("expired app capability token")
)

type Claims struct {
	Issuer             string   `json:"iss"`
	Audience           string   `json:"aud"`
	Subject            string   `json:"sub"`
	AppID              string   `json:"app_id"`
	AppInstallationID  string   `json:"app_installation_id,omitempty"`
	WorkspaceID        string   `json:"workspace_id"`
	MemberID           string   `json:"member_id"`
	AgentID            string   `json:"agent_id,omitempty"`
	AgentMemberID      string   `json:"agent_member_id,omitempty"`
	AuthorizerMemberID string   `json:"authorizer_member_id,omitempty"`
	SessionID          string   `json:"session_id,omitempty"`
	PolicyVersion      int64    `json:"policy_version,omitempty"`
	Scopes             []string `json:"scopes"`
	TokenID            string   `json:"jti"`
	RequestID          string   `json:"request_id,omitempty"`
	IssuedAtUnix       int64    `json:"iat"`
	ExpiresAtUnix      int64    `json:"exp"`
}

type Signer struct {
	key []byte
}

func NewSigner(key []byte) (*Signer, error) {
	if len(key) == 0 {
		return nil, fmt.Errorf("app capability signing key is required")
	}
	copied := make([]byte, len(key))
	copy(copied, key)
	return &Signer{key: copied}, nil
}

func NewEphemeralSigner() (*Signer, error) {
	key := make([]byte, 32)
	if _, err := rand.Read(key); err != nil {
		return nil, fmt.Errorf("generate app capability signing key: %w", err)
	}
	return NewSigner(key)
}

func Issuer() string {
	return issuer
}

func Audience() string {
	return audience
}

func (signer *Signer) Sign(claims Claims) (string, error) {
	if signer == nil || len(signer.key) == 0 {
		return "", fmt.Errorf("app capability signer is not configured")
	}
	normalized, err := normalizeSigningClaims(claims)
	if err != nil {
		return "", err
	}

	headerJSON, err := json.Marshal(map[string]string{
		"alg": "HS256",
		"typ": "JWT",
	})
	if err != nil {
		return "", fmt.Errorf("encode app capability header: %w", err)
	}
	payloadJSON, err := json.Marshal(normalized)
	if err != nil {
		return "", fmt.Errorf("encode app capability claims: %w", err)
	}

	headerSegment := base64.RawURLEncoding.EncodeToString(headerJSON)
	payloadSegment := base64.RawURLEncoding.EncodeToString(payloadJSON)
	signingInput := headerSegment + "." + payloadSegment
	signature := signer.signature(signingInput)
	return signingInput + "." + base64.RawURLEncoding.EncodeToString(signature), nil
}

func (signer *Signer) Verify(token string) (Claims, error) {
	if signer == nil || len(signer.key) == 0 {
		return Claims{}, fmt.Errorf("app capability signer is not configured")
	}

	parts := strings.Split(strings.TrimSpace(token), ".")
	if len(parts) != 3 {
		return Claims{}, ErrInvalidToken
	}

	signingInput := parts[0] + "." + parts[1]
	signature, err := base64.RawURLEncoding.DecodeString(parts[2])
	if err != nil {
		return Claims{}, ErrInvalidToken
	}
	expected := signer.signature(signingInput)
	if !hmac.Equal(signature, expected) {
		return Claims{}, ErrInvalidToken
	}

	headerJSON, err := base64.RawURLEncoding.DecodeString(parts[0])
	if err != nil {
		return Claims{}, ErrInvalidToken
	}
	var header struct {
		Algorithm string `json:"alg"`
		Type      string `json:"typ"`
	}
	if err := json.Unmarshal(headerJSON, &header); err != nil {
		return Claims{}, ErrInvalidToken
	}
	if header.Algorithm != "HS256" || header.Type != "JWT" {
		return Claims{}, ErrInvalidToken
	}

	payloadJSON, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return Claims{}, ErrInvalidToken
	}
	var claims Claims
	if err := json.Unmarshal(payloadJSON, &claims); err != nil {
		return Claims{}, ErrInvalidToken
	}
	normalized, err := normalizeVerifiedClaims(claims)
	if err != nil {
		return Claims{}, ErrInvalidToken
	}
	if time.Now().UTC().Unix() >= normalized.ExpiresAtUnix {
		return Claims{}, ErrExpiredToken
	}
	return normalized, nil
}

func (signer *Signer) signature(signingInput string) []byte {
	mac := hmac.New(sha256.New, signer.key)
	_, _ = mac.Write([]byte(signingInput))
	return mac.Sum(nil)
}

func normalizeSigningClaims(claims Claims) (Claims, error) {
	claims.Issuer = strings.TrimSpace(claims.Issuer)
	if claims.Issuer == "" {
		claims.Issuer = issuer
	}
	claims.Audience = strings.TrimSpace(claims.Audience)
	if claims.Audience == "" {
		claims.Audience = audience
	}
	return validateClaims(normalizeClaimsFields(claims))
}

func normalizeVerifiedClaims(claims Claims) (Claims, error) {
	claims.Issuer = strings.TrimSpace(claims.Issuer)
	claims.Audience = strings.TrimSpace(claims.Audience)
	return validateClaims(normalizeClaimsFields(claims))
}

func normalizeClaimsFields(claims Claims) Claims {
	claims.Subject = strings.TrimSpace(claims.Subject)
	claims.AppID = strings.TrimSpace(claims.AppID)
	claims.AppInstallationID = strings.TrimSpace(claims.AppInstallationID)
	claims.WorkspaceID = strings.TrimSpace(claims.WorkspaceID)
	claims.MemberID = strings.TrimSpace(claims.MemberID)
	claims.AgentID = strings.TrimSpace(claims.AgentID)
	claims.AgentMemberID = strings.TrimSpace(claims.AgentMemberID)
	claims.AuthorizerMemberID = strings.TrimSpace(claims.AuthorizerMemberID)
	claims.SessionID = strings.TrimSpace(claims.SessionID)
	claims.TokenID = strings.TrimSpace(claims.TokenID)
	claims.RequestID = strings.TrimSpace(claims.RequestID)
	claims.Scopes = normalizeScopes(claims.Scopes)
	return claims
}

func validateClaims(claims Claims) (Claims, error) {
	if claims.Issuer != issuer || claims.Audience != audience {
		return Claims{}, ErrInvalidToken
	}
	if claims.Subject == "" || claims.AppID == "" || claims.WorkspaceID == "" || claims.MemberID == "" || claims.TokenID == "" {
		return Claims{}, ErrInvalidToken
	}
	if claims.IssuedAtUnix <= 0 || claims.ExpiresAtUnix <= claims.IssuedAtUnix {
		return Claims{}, ErrInvalidToken
	}
	return claims, nil
}

func normalizeScopes(scopes []string) []string {
	if len(scopes) == 0 {
		return nil
	}

	seen := make(map[string]struct{}, len(scopes))
	normalized := make([]string, 0, len(scopes))
	for _, scope := range scopes {
		trimmed := strings.TrimSpace(scope)
		if trimmed == "" {
			continue
		}
		if _, exists := seen[trimmed]; exists {
			continue
		}
		seen[trimmed] = struct{}{}
		normalized = append(normalized, trimmed)
	}
	if len(normalized) == 0 {
		return nil
	}
	return normalized
}