package store

import (
	"context"
	"path/filepath"
	"testing"
	"time"
)

func selfUsageFixture(t *testing.T, s *DB) time.Time {
	t.Helper()
	ctx := context.Background()
	must := func(err error) {
		t.Helper()
		if err != nil {
			t.Fatal(err)
		}
	}
	must(s.PutOrg(ctx, Org{ID: "o", Name: "Acme"}))
	must(s.PutTeam(ctx, Team{ID: "t", OrgID: "o", Name: "ml"}))
	must(s.PutProject(ctx, Project{ID: "p", TeamID: "t", Name: "training"}))
	must(s.PutProject(ctx, Project{ID: "q", TeamID: "t", Name: "eval"}))
	now := time.Now()
	u := func(user, project string, cost int64) Usage {
		return Usage{At: now, KeyID: "session:" + user, UserID: user, ProjectID: project,
			Model: "m", ServedModel: "m", PromptTokens: 10, CompletionTokens: 5, CostMicros: cost, Status: 200}
	}
	must(s.WriteUsage(ctx, []Usage{
		u("ada", "p", 100), u("ada", "p", 100), u("ada", "", 30), u("bob", "p", 1000), u("bob", "q", 7),
	}))
	return now
}

func checkSelfUsage(t *testing.T, s *DB) {
	now := selfUsageFixture(t, s)
	ctx := context.Background()
	since := now.Add(-48 * time.Hour)

	places, buckets, err := s.UsageOfUser(ctx, "ada", since, 24*time.Hour)
	if err != nil {
		t.Fatal(err)
	}
	got := map[string]PlaceUsage{}
	for _, pl := range places {
		got[pl.ProjectID] = pl
	}
	// Ada's own, and none of Bob's.
	if len(places) != 2 || got["p"].Requests != 2 || got["p"].Path != "Acme / ml / training" ||
		got[""].CostMicros != 30 || got["p"].Tokens != 30 {
		t.Errorf("places = %+v", places)
	}
	var total int64
	for _, b := range buckets {
		total += b.Requests
	}
	if total != 3 || len(buckets) < 2 {
		t.Errorf("%d buckets carrying %d requests", len(buckets), total)
	}

	for _, c := range []struct {
		l    Limits
		want int64
	}{
		{Limits{Scope: ScopeUser, ScopeID: "ada"}, 230},
		{Limits{Scope: ScopeProject, ScopeID: "p"}, 1200},
		{Limits{Scope: ScopeTeam, ScopeID: "t"}, 1207},
		{Limits{Scope: ScopeOrg, ScopeID: "o"}, 1207},
		{Limits{Scope: ScopeProjectMember, ScopeID: "p/ada"}, 200},
		{Limits{Scope: ScopeTeamMember, ScopeID: "t/bob"}, 1007},
	} {
		n, err := s.SpendOfPool(ctx, c.l, since)
		if err != nil {
			t.Fatal(err)
		}
		if n != c.want {
			t.Errorf("%s %s spent %d, want %d", c.l.Scope, c.l.ScopeID, n, c.want)
		}
	}
}

func TestAPersonSeesTheirOwnUse(t *testing.T) {
	s, err := OpenSQLite(context.Background(), filepath.Join(t.TempDir(), "u.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = s.Close() })
	checkSelfUsage(t, s)
}

func TestPostgresAPersonSeesTheirOwnUse(t *testing.T) {
	checkSelfUsage(t, openPG(t))
}