package llm

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	ck "gitea.maximumdirect.net/eric/ConfigKit"
	"gitea.maximumdirect.net/eric/audita/internal/core/config"
	"gitea.maximumdirect.net/eric/audita/internal/framework/outcomes"
	"gitea.maximumdirect.net/eric/audita/internal/framework/responseschema"
	"gitea.maximumdirect.net/eric/audita/internal/framework/structuredoutput"
	"gitea.maximumdirect.net/eric/audita/internal/prompts"
	pk "gitea.maximumdirect.net/eric/promptkit"
)

func TestRolesShareBackendCapacityAndReleaseAfterCancellation(t *testing.T) {
	t.Setenv("INHERITED_TEST_KEY", "synthetic-inherited-runtime-secret")
	entered := make(chan struct{}, 3)
	release := make(chan struct{})
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if got := r.Header.Get("Authorization"); got != "Bearer synthetic-inherited-runtime-secret" {
			t.Errorf("inherited credential header: %q", got)
		}
		entered <- struct{}{}
		<-release
		content, _ := json.Marshal(`{"corrections":[]}`)
		_, _ = fmt.Fprintf(w, `{"choices":[{"message":{"content":%s}}]}`, content)
	}))
	defer server.Close()
	dir := t.TempDir()
	if err := os.Mkdir(filepath.Join(dir, "profiles"), 0700); err != nil {
		t.Fatal(err)
	}
	main := filepath.Join(dir, "config.yml")
	body := fmt.Sprintf("version: 2\nbackends:\n  shared-base:\n    endpoint: %s/v1\n    credential: INHERITED_TEST_KEY\n    concurrency: 1\n    queue_capacity: 0\n  shared:\n    extends: shared-base\nllm:\n  proposal_profile: proposal\n  validation_profile: validation\n  transient_retries: 0\n  repair_attempts: 0\n", server.URL)
	if err := os.WriteFile(main, []byte(body), 0600); err != nil {
		t.Fatal(err)
	}
	for _, name := range []string{"proposal", "validation"} {
		profile := "profiles:\n  " + name + ":\n    backend: shared\n    model: local\n"
		if err := os.WriteFile(filepath.Join(dir, "profiles", name+".yaml"), []byte(profile), 0600); err != nil {
			t.Fatal(err)
		}
	}
	loaded, err := config.Load(ck.File(main), "")
	if err != nil {
		t.Fatal(err)
	}
	runtime, err := NewRuntime(context.Background(), loaded.Config, main, server.Client())
	if err != nil {
		t.Fatal(err)
	}
	proposal, err := runtime.NewRoleClient("proposal", CallHooks{})
	if err != nil {
		t.Fatal(err)
	}
	validation, err := runtime.NewRoleClient("validation", CallHooks{})
	if err != nil {
		t.Fatal(err)
	}
	request := structuredoutput.Request{PromptID: prompts.PromptIDModuleGlossaryProposal, PromptVersion: prompts.VersionV1, Variables: map[string]string{"TranscriptDescriptionBlock": "", "GlossaryJSON": "[]", "SectionJSON": "[]"}, Schema: responseschema.CorrectionSetKey}
	firstDone := make(chan error, 1)
	go func() {
		var output map[string]any
		_, callErr := proposal.CompleteStructured(context.Background(), request, &output)
		firstDone <- callErr
	}()
	select {
	case <-entered:
	case <-time.After(2 * time.Second):
		t.Fatal("first call did not enter provider")
	}
	var output map[string]any
	capacityResponse, err := validation.CompleteStructured(context.Background(), request, &output)
	if !errors.Is(err, pk.ErrCapacityExceeded) {
		t.Fatalf("other role did not share capacity: %v", err)
	}
	var capacityFailure *structuredoutput.Failure
	if !errors.As(err, &capacityFailure) || capacityFailure.Class != structuredoutput.FailureRecoverable || capacityFailure.Reason != outcomes.ReasonCapacityExceeded || !capacityResponse.Execution.Available || capacityResponse.Execution.ProviderAttempts != 0 {
		t.Fatalf("capacity outcome: failure=%#v counts=%+v", capacityFailure, capacityResponse.Execution)
	}
	select {
	case <-entered:
		t.Fatal("second role reached provider despite full backend")
	default:
	}
	close(release)
	if err := <-firstDone; err != nil {
		t.Fatal(err)
	}
	_, err = validation.CompleteStructured(context.Background(), request, &output)
	if err != nil {
		t.Fatalf("capacity was not released: %v", err)
	}
	for len(entered) > 0 {
		<-entered
	}
	release = make(chan struct{})
	if err := os.WriteFile(main, []byte(strings.Replace(body, "queue_capacity: 0", "queue_capacity: 1", 1)), 0600); err != nil {
		t.Fatal(err)
	}
	loaded, err = config.Load(ck.File(main), "")
	if err != nil {
		t.Fatal(err)
	}
	runtime, err = NewRuntime(context.Background(), loaded.Config, main, server.Client())
	if err != nil {
		t.Fatal(err)
	}
	proposal, err = runtime.NewRoleClient("proposal", CallHooks{})
	if err != nil {
		t.Fatal(err)
	}
	validation, err = runtime.NewRoleClient("validation", CallHooks{})
	if err != nil {
		t.Fatal(err)
	}
	go func() {
		var out map[string]any
		_, callErr := proposal.CompleteStructured(context.Background(), request, &out)
		firstDone <- callErr
	}()
	select {
	case <-entered:
	case <-time.After(2 * time.Second):
		t.Fatal("first queued-capacity call did not dispatch")
	}
	queuedCtx, cancelQueued := context.WithTimeout(context.Background(), 30*time.Millisecond)
	defer cancelQueued()
	_, err = validation.CompleteStructured(queuedCtx, request, &output)
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("queued call did not honor deadline: %v", err)
	}
	if !errors.As(err, &capacityFailure) || capacityFailure.Class != structuredoutput.FailureCanceled || capacityFailure.Reason != outcomes.ReasonCanceled {
		t.Fatalf("queued parent deadline classification: %#v", capacityFailure)
	}
	close(release)
	if err := <-firstDone; err != nil {
		t.Fatal(err)
	}
	_, err = validation.CompleteStructured(context.Background(), request, &output)
	if err != nil {
		t.Fatalf("queued cancellation leaked capacity: %v", err)
	}
}

func TestDistinctBackendsHaveIndependentCapacity(t *testing.T) {
	entered := make(chan struct{}, 2)
	release := make(chan struct{})
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		entered <- struct{}{}
		<-release
		schema := map[string]any{}
		_ = json.NewDecoder(r.Body).Decode(&schema)
		content, _ := json.Marshal(`{"corrections":[]}`)
		_, _ = fmt.Fprintf(w, `{"choices":[{"message":{"content":%s}}]}`, content)
	}))
	defer server.Close()
	dir := t.TempDir()
	if err := os.Mkdir(filepath.Join(dir, "profiles"), 0700); err != nil {
		t.Fatal(err)
	}
	main := filepath.Join(dir, "config.yml")
	body := fmt.Sprintf("version: 2\nbackends:\n  first:\n    endpoint: %s/v1\n    concurrency: 1\n    queue_capacity: 0\n  second:\n    endpoint: %s/v1\n    concurrency: 1\n    queue_capacity: 0\nllm:\n  proposal_profile: proposal\n  validation_profile: validation\n  transient_retries: 0\n  repair_attempts: 0\n", server.URL, server.URL)
	if err := os.WriteFile(main, []byte(body), 0600); err != nil {
		t.Fatal(err)
	}
	for name, backend := range map[string]string{"proposal": "first", "validation": "second"} {
		if err := os.WriteFile(filepath.Join(dir, "profiles", name+".yaml"), []byte("profiles:\n  "+name+":\n    backend: "+backend+"\n    model: local\n"), 0600); err != nil {
			t.Fatal(err)
		}
	}
	loaded, err := config.Load(ck.File(main), "")
	if err != nil {
		t.Fatal(err)
	}
	runtime, err := NewRuntime(context.Background(), loaded.Config, main, server.Client())
	if err != nil {
		t.Fatal(err)
	}
	proposal, err := runtime.NewRoleClient("proposal", CallHooks{})
	if err != nil {
		t.Fatal(err)
	}
	validation, err := runtime.NewRoleClient("validation", CallHooks{})
	if err != nil {
		t.Fatal(err)
	}
	request := structuredoutput.Request{PromptID: prompts.PromptIDModuleGlossaryProposal, PromptVersion: prompts.VersionV1, Variables: map[string]string{"TranscriptDescriptionBlock": "", "GlossaryJSON": "[]", "SectionJSON": "[]"}, Schema: responseschema.CorrectionSetKey}
	done := make(chan error, 2)
	go func() {
		var out map[string]any
		_, e := proposal.CompleteStructured(context.Background(), request, &out)
		done <- e
	}()
	select {
	case <-entered:
	case <-time.After(2 * time.Second):
		t.Fatal("first backend did not dispatch")
	}
	go func() {
		var out map[string]any
		_, e := validation.CompleteStructured(context.Background(), request, &out)
		done <- e
	}()
	select {
	case <-entered:
	case <-time.After(2 * time.Second):
		t.Fatal("second backend was blocked by first")
	}
	close(release)
	for range 2 {
		if err := <-done; err != nil {
			t.Fatal(err)
		}
	}
}