package main

import (
	"context"
	"testing"

	"git.cin.su/m/Assistant/contracts/capability"
	"git.cin.su/m/Assistant/contracts/cognitive"
)

type admissionTestProvider struct{ calls int }

func (p *admissionTestProvider) Advise(context.Context, cognitive.AdvisoryRequest) (cognitive.AdvisoryResponse, error) {
	p.calls++
	return cognitive.AdvisoryResponse{FinishReason: cognitive.FinishCompleted}, nil
}

type admissionTestClient struct {
	decision capability.AdmissionDecision
	releases int
}

func (c *admissionTestClient) Admit(context.Context, string, capability.AdmissionRequest) (capability.AdmissionDecision, capability.AdmissionLease, error) {
	return c.decision, capability.AdmissionLease{ID: "lease-1", PolicyID: "policy-1"}, nil
}
func (c *admissionTestClient) Release(context.Context, capability.AdmissionLease) error {
	c.releases++
	return nil
}

func TestInvokeAdmittedDoesNotCallProviderWhenRejected(t *testing.T) {
	provider := &admissionTestProvider{}
	client := &admissionTestClient{decision: capability.AdmissionDecision{State: capability.AdmissionRejected, Reason: capability.AdmissionReasonConcurrency, PolicyID: "policy-1"}}
	gateway := cognitiveGateway{admission: client}
	response, err := gateway.invokeAdmitted(context.Background(), "request-1", provider, cognitive.AdvisoryRequest{Message: "hello"}, 100)
	if err != nil || response.FinishReason != cognitive.FinishRateLimited || provider.calls != 0 || client.releases != 0 {
		t.Fatalf("rejected admission leaked provider call: response=%#v err=%v calls=%d releases=%d", response, err, provider.calls, client.releases)
	}
}

func TestInvokeAdmittedReleasesDurableLeaseAfterProvider(t *testing.T) {
	provider := &admissionTestProvider{}
	client := &admissionTestClient{decision: capability.AdmissionDecision{State: capability.AdmissionAdmitted, Reason: capability.AdmissionReasonNone, PolicyID: "policy-1"}}
	gateway := cognitiveGateway{admission: client}
	response, err := gateway.invokeAdmitted(context.Background(), "request-1", provider, cognitive.AdvisoryRequest{Message: "hello"}, 100)
	if err != nil || response.FinishReason != cognitive.FinishCompleted || provider.calls != 1 || client.releases != 1 {
		t.Fatalf("admitted lease was not released: response=%#v err=%v calls=%d releases=%d", response, err, provider.calls, client.releases)
	}
}