package main

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

// testAccount is the fake ChatGPT account ID in the fake tokens.
const testAccount = "account-test"

func TestLoginWritesTheTokenFile(t *testing.T) {
	server := newFakeAuth(t)
	server.pending = 2
	auth := newTestAuth(t, server)

	var out strings.Builder
	err := auth.Login(t.Context(), &out)
	require.NoError(t, err)

	require.Contains(t, out.String(), server.URL+"/codex/device")
	require.Contains(t, out.String(), "CODE-1234")
	require.Equal(t, 3, server.polls)

	form := server.exchange
	require.Equal(t, "authorization_code", form.Get("grant_type"))
	require.Equal(t, "auth-code", form.Get("code"))
	require.Equal(t, "verifier", form.Get("code_verifier"))
	require.Equal(t, codexClientID, form.Get("client_id"))
	require.Equal(t, server.URL+"/deviceauth/callback", form.Get("redirect_uri"))

	tokens := readTokenFile(t, auth.path)
	require.Equal(t, "refresh-1", tokens.RefreshToken)
	require.Equal(t, testAccount, tokens.AccountID)
	require.Equal(t, server.access[0], tokens.AccessToken)
	require.NotContains(t, out.String(), tokens.AccessToken)
	require.NotContains(t, out.String(), tokens.RefreshToken)

	info, err := os.Stat(auth.path)
	require.NoError(t, err)
	require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
}

func TestLoginFailsWhenTheCodeIsDenied(t *testing.T) {
	server := newFakeAuth(t)
	server.deny = true
	auth := newTestAuth(t, server)

	err := auth.Login(t.Context(), io.Discard)
	require.ErrorContains(t, err, "HTTP 400")
	require.NoFileExists(t, auth.path)
}

func TestAuthorizeUsesAFreshTokenWithoutRefresh(t *testing.T) {
	server := newFakeAuth(t)
	auth := newTestAuth(t, server)
	writeTestTokens(t, auth, server.token(time.Hour), "refresh-0")

	header := http.Header{}
	err := auth.Authorize(t.Context(), header)
	require.NoError(t, err)
	require.Equal(t, "Bearer "+server.access[0], header.Get("Authorization"))
	require.Equal(t, testAccount, header.Get("ChatGPT-Account-Id"))
	require.Equal(t, 0, server.refreshes)
}

func TestAuthorizeRefreshesNearTheExpiry(t *testing.T) {
	testCases := []struct {
		name        string
		newRefresh  string
		wantRefresh string
	}{
		{
			name:        "new refresh token kept",
			newRefresh:  "refresh-new",
			wantRefresh: "refresh-new",
		},
		{
			name:        "old refresh token kept without a new one",
			newRefresh:  "",
			wantRefresh: "refresh-0",
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			server := newFakeAuth(t)
			server.newRefresh = tc.newRefresh
			auth := newTestAuth(t, server)
			writeTestTokens(t, auth, server.token(59*time.Second), "refresh-0")

			header := http.Header{}
			err := auth.Authorize(t.Context(), header)
			require.NoError(t, err)
			require.Equal(t, 1, server.refreshes)
			require.Equal(t, "refresh-0", server.refreshed)

			tokens := readTokenFile(t, auth.path)
			require.Equal(t, server.access[1], tokens.AccessToken)
			require.Equal(t, tc.wantRefresh, tokens.RefreshToken)
			require.Equal(t, testAccount, tokens.AccountID)
			require.Equal(t, "Bearer "+server.access[1], header.Get("Authorization"))

			err = auth.Authorize(t.Context(), header)
			require.NoError(t, err)
			require.Equal(t, 1, server.refreshes)
		})
	}
}

func TestAuthorizeRunsOneRefreshAtATime(t *testing.T) {
	server := newFakeAuth(t)
	server.newRefresh = "refresh-new"
	auth := newTestAuth(t, server)
	writeTestTokens(t, auth, server.token(time.Second), "refresh-0")

	var group sync.WaitGroup
	errs := make(chan error, 8)
	for range 8 {
		group.Go(func() {
			errs <- auth.Authorize(t.Context(), http.Header{})
		})
	}

	group.Wait()
	close(errs)
	for err := range errs {
		require.NoError(t, err)
	}

	require.Equal(t, 1, server.refreshes)
}

func TestAuthorizeNeedsALogin(t *testing.T) {
	server := newFakeAuth(t)
	auth := newTestAuth(t, server)

	err := auth.Authorize(t.Context(), http.Header{})
	require.ErrorIs(t, err, errLoginNeeded)

	server.failRefresh = true
	writeTestTokens(t, auth, server.token(-time.Minute), "refresh-0")
	before, err := os.ReadFile(auth.path)
	require.NoError(t, err)

	err = auth.Authorize(t.Context(), http.Header{})
	require.ErrorIs(t, err, errLoginNeeded)
	require.NotContains(t, err.Error(), "refresh-0")

	after, err := os.ReadFile(auth.path)
	require.NoError(t, err)
	require.Equal(t, before, after)
}

func TestForcedRefreshLetsTheNextCallWork(t *testing.T) {
	server := newFakeAuth(t)
	auth := newTestAuth(t, server)
	writeTestTokens(t, auth, server.token(-time.Hour), "refresh-0")

	responses := newModelFake(t)
	client := newModelClient(responses.URL, auth, "agent-key", filepath.Join(t.TempDir(), "usage.jsonl"))
	_, err := client.Create(t.Context(), modelTestRequest())
	require.NoError(t, err)

	require.Equal(t, 1, server.refreshes)
	call := responses.Call(0)
	require.Equal(t, "Bearer "+server.access[1], call.header.Get("Authorization"))
	require.Equal(t, testAccount, call.header.Get("ChatGPT-Account-Id"))
	require.Equal(t, server.access[1], readTokenFile(t, auth.path).AccessToken)
}

func TestModelEndpointFollowsTheConfig(t *testing.T) {
	dir := t.TempDir()
	config := defaultConfig()

	t.Setenv("OPENAI_API_KEY", modelTestKey)
	address, auth, err := modelEndpoint(config, dir)
	require.NoError(t, err)
	require.Equal(t, modelChatGPTURL, address)
	require.IsType(t, &chatgptAuth{}, auth)
	require.Equal(t, filepath.Join(dir, tokenFileName), auth.(*chatgptAuth).path)

	config.Endpoint = endpointAPI
	address, auth, err = modelEndpoint(config, dir)
	require.NoError(t, err)
	require.Equal(t, modelAPIURL, address)
	require.Equal(t, ModelAPIKey(modelTestKey), auth)

	t.Setenv("OPENAI_API_KEY", "")
	_, _, err = modelEndpoint(config, dir)
	require.ErrorContains(t, err, "OPENAI_API_KEY")
}

// fakeAuth is an httptest login service. The device token endpoint answers
// 403 pending times before it grants the code. Each token answer has a new
// access token, kept in access.
type fakeAuth struct {
	*httptest.Server
	mutex       sync.Mutex
	pending     int
	deny        bool
	polls       int
	exchange    url.Values
	refreshes   int
	refreshed   string
	newRefresh  string
	failRefresh bool
	access      []string
}

func newFakeAuth(t *testing.T) *fakeAuth {
	t.Helper()

	f := &fakeAuth{}
	mux := http.NewServeMux()
	mux.HandleFunc("POST /api/accounts/deviceauth/usercode", f.userCode)
	mux.HandleFunc("POST /api/accounts/deviceauth/token", f.deviceToken)
	mux.HandleFunc("POST /oauth/token", f.oauthToken)
	f.Server = httptest.NewServer(mux)
	t.Cleanup(f.Close)

	return f
}

func (f *fakeAuth) userCode(w http.ResponseWriter, r *http.Request) {
	var body map[string]string
	err := json.NewDecoder(r.Body).Decode(&body)
	if err != nil || body["client_id"] != codexClientID {
		http.Error(w, "bad client", http.StatusBadRequest)
		return
	}

	writeJSON(w, map[string]any{
		"device_auth_id": "device-1",
		"user_code":      "CODE-1234",
		"interval":       "0",
	})
}

func (f *fakeAuth) deviceToken(w http.ResponseWriter, r *http.Request) {
	f.mutex.Lock()
	defer f.mutex.Unlock()

	var body map[string]string
	err := json.NewDecoder(r.Body).Decode(&body)
	if err != nil || body["device_auth_id"] != "device-1" || body["user_code"] != "CODE-1234" {
		http.Error(w, "bad device", http.StatusBadRequest)
		return
	}

	f.polls++
	if f.deny {
		http.Error(w, "denied", http.StatusBadRequest)
		return
	}

	if f.polls <= f.pending {
		http.Error(w, "pending", http.StatusForbidden)
		return
	}

	writeJSON(w, map[string]any{
		"authorization_code": "auth-code",
		"code_challenge":     "challenge",
		"code_verifier":      "verifier",
	})
}

func (f *fakeAuth) oauthToken(w http.ResponseWriter, r *http.Request) {
	f.mutex.Lock()
	defer f.mutex.Unlock()

	if r.Header.Get("Content-Type") == "application/x-www-form-urlencoded" {
		err := r.ParseForm()
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		f.exchange = r.PostForm
		writeJSON(w, map[string]any{
			"access_token":  f.newToken(time.Hour),
			"refresh_token": "refresh-1",
			"id_token":      f.newToken(time.Hour),
		})

		return
	}

	var body map[string]string
	err := json.NewDecoder(r.Body).Decode(&body)
	if err != nil || body["grant_type"] != "refresh_token" || body["client_id"] != codexClientID {
		http.Error(w, "bad refresh", http.StatusBadRequest)
		return
	}

	f.refreshes++
	f.refreshed = body["refresh_token"]
	if f.failRefresh {
		http.Error(w, `{"error": "invalid_grant", "token": "`+body["refresh_token"]+`"}`, http.StatusUnauthorized)
		return
	}

	answer := map[string]any{
		"access_token": f.newToken(time.Hour),
		"id_token":     f.newToken(time.Hour),
	}

	if f.newRefresh != "" {
		answer["refresh_token"] = f.newRefresh
	}

	writeJSON(w, answer)
}

// token returns a new access token that expires after ttl, for a token file
// that a test writes.
func (f *fakeAuth) token(ttl time.Duration) string {
	f.mutex.Lock()
	defer f.mutex.Unlock()

	return f.newToken(ttl)
}

// newToken returns a new fake JWT with the expiry and the account claim, and
// keeps it in access. The caller holds the mutex.
func (f *fakeAuth) newToken(ttl time.Duration) string {
	claims, err := json.Marshal(map[string]any{
		"exp": time.Now().Add(ttl).Unix(),
		authClaim: map[string]any{
			"chatgpt_account_id": testAccount,
		},
		"serial": len(f.access),
	})
	if err != nil {
		panic(err)
	}

	encode := base64.RawURLEncoding.EncodeToString
	token := encode([]byte(`{"alg":"none"}`)) + "." + encode(claims) + "." + encode([]byte("fake"))
	f.access = append(f.access, token)

	return token
}

func newTestAuth(t *testing.T, server *fakeAuth) *chatgptAuth {
	t.Helper()

	return newChatGPTAuth(filepath.Join(t.TempDir(), tokenFileName), server.URL)
}

func writeTestTokens(t *testing.T, auth *chatgptAuth, access, refresh string) {
	t.Helper()

	err := auth.write(tokenFile{
		AccessToken:  access,
		RefreshToken: refresh,
		IDToken:      access,
		AccountID:    testAccount,
	})
	require.NoError(t, err)
}

func readTokenFile(t *testing.T, path string) tokenFile {
	t.Helper()

	data, err := os.ReadFile(path)
	require.NoError(t, err)

	var tokens tokenFile
	require.NoError(t, json.Unmarshal(data, &tokens))

	return tokens
}

func writeJSON(w http.ResponseWriter, value any) {
	w.Header().Set("Content-Type", "application/json")
	err := json.NewEncoder(w).Encode(value)
	if err != nil {
		panic(fmt.Sprint("encode answer: ", err))
	}
}