package authproxy

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"slices"
	"strconv"
	"strings"
	"time"
)

// idToken holds the claims we care about from Google's OpenID Connect
// id_token. The token arrives directly from Google's token endpoint over
// TLS, so per the OIDC spec its signature does not need separate
// verification; we still validate issuer, audience, and expiry.
type idToken struct {
	Iss           string   `json:"iss"`
	Sub           string   `json:"sub"`
	Aud           flexAud  `json:"aud"`
	Exp           int64    `json:"exp"`
	Email         string   `json:"email"`
	EmailVerified flexBool `json:"email_verified"`
	Hd            string   `json:"hd"`
	Name          string   `json:"name"`
}

func (s *Server) exchangeCode(ctx context.Context, code string) (*idToken, error) {
	if code == "" {
		return nil, fmt.Errorf("missing code parameter")
	}
	form := url.Values{
		"code":          {code},
		"client_id":     {s.cfg.ClientID},
		"client_secret": {s.cfg.ClientSecret},
		"redirect_uri":  {s.redirectURI()},
		"grant_type":    {"authorization_code"},
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, s.cfg.TokenURL, strings.NewReader(form.Encode()))
	if err != nil {
		return nil, err
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")

	resp, err := s.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("token endpoint: %w", err)
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return nil, fmt.Errorf("token endpoint read: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("token endpoint returned %d: %s", resp.StatusCode, truncate(string(body), 200))
	}

	var tr struct {
		IDToken string `json:"id_token"`
	}
	if err := json.Unmarshal(body, &tr); err != nil {
		return nil, fmt.Errorf("token endpoint response: %w", err)
	}
	if tr.IDToken == "" {
		return nil, fmt.Errorf("token endpoint response missing id_token")
	}
	return parseIDToken(tr.IDToken)
}

func parseIDToken(raw string) (*idToken, error) {
	parts := strings.Split(raw, ".")
	if len(parts) != 3 {
		return nil, fmt.Errorf("id_token is not a JWT")
	}
	payload, err := base64.RawURLEncoding.DecodeString(parts[1])
	if err != nil {
		return nil, fmt.Errorf("id_token payload: %w", err)
	}
	var tok idToken
	if err := json.Unmarshal(payload, &tok); err != nil {
		return nil, fmt.Errorf("id_token claims: %w", err)
	}
	return &tok, nil
}

func (s *Server) validateIDToken(t *idToken) error {
	if t.Iss != "https://accounts.google.com" && t.Iss != "accounts.google.com" {
		return fmt.Errorf("unexpected issuer %q", t.Iss)
	}
	if !t.Aud.contains(s.cfg.ClientID) {
		return fmt.Errorf("audience mismatch")
	}
	if time.Now().Unix() > t.Exp {
		return fmt.Errorf("token expired")
	}
	if t.Email == "" {
		return fmt.Errorf("no email claim")
	}
	if !bool(t.EmailVerified) {
		return fmt.Errorf("email %s is not verified", t.Email)
	}
	return nil
}

// flexAud accepts the JWT aud claim as either a string or an array.
type flexAud []string

func (a *flexAud) UnmarshalJSON(b []byte) error {
	var single string
	if err := json.Unmarshal(b, &single); err == nil {
		*a = flexAud{single}
		return nil
	}
	var many []string
	if err := json.Unmarshal(b, &many); err != nil {
		return err
	}
	*a = flexAud(many)
	return nil
}

func (a flexAud) contains(v string) bool {
	return slices.Contains(a, v)
}

// flexBool accepts true, "true", false, or "false" — Google has historically
// been inconsistent about the email_verified type.
type flexBool bool

func (b *flexBool) UnmarshalJSON(data []byte) error {
	s := strings.Trim(string(data), `"`)
	v, err := strconv.ParseBool(s)
	if err != nil {
		return fmt.Errorf("invalid bool value %s", data)
	}
	*b = flexBool(v)
	return nil
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "..."
}