package main

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

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

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

// ChatGPTFake 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 ChatGPTFake 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
}

// token returns a new access token that expires after ttl, for a token file
// that a test writes.
func (f *ChatGPTFake) 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 *ChatGPTFake) newToken(ttl time.Duration) string {
	claims, err := json.Marshal(map[string]any{
		"exp": time.Now().Add(ttl).Unix(),
		chatgptAuthClaim: map[string]any{
			"chatgpt_account_id": chatgptTestAccount,
		},
		"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 (f *ChatGPTFake) 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"] != chatgptClientID {
		http.Error(w, "bad client", http.StatusBadRequest)

		return
	}

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

func (f *ChatGPTFake) 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
	}

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

func (f *ChatGPTFake) 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
		chatgptWriteJSON(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"] != chatgptClientID {
		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
	}

	chatgptWriteJSON(w, answer)
}

func newChatGPTFake(t *testing.T) *ChatGPTFake {
	t.Helper()

	f := &ChatGPTFake{}
	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 newTestChatGPTAuth(t *testing.T, server *ChatGPTFake) *ChatGPTAuth {
	t.Helper()

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

func chatgptWriteTestTokens(t *testing.T, auth *ChatGPTAuth, access, refresh string) {
	t.Helper()

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

func chatgptWriteJSON(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))
	}
}