package llm

import (
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/promptkit"
)

func newObservedHTTPClient(t *testing.T, handler http.HandlerFunc, retries int, transport http.RoundTripper) *PromptKitClient {
	t.Helper()
	server := httptest.NewServer(handler)
	t.Cleanup(server.Close)
	client, err := NewPromptKitClient(PromptKitClientConfig{
		Assets: newTestPromptKitAssets(t), TransientRetries: &retries,
		HTTPClient: &http.Client{Transport: transport},
		EngineOptions: []promptkit.Option{promptkit.WithProfiles(promptkit.OpenAICompatibleProfile(promptkit.OpenAICompatibleProfileConfig{
			ID: "default-profile", Endpoint: server.URL + "/v1", Model: "test-model",
		}))},
	})
	if err != nil {
		t.Fatal(err)
	}
	return client
}

func observedRequest() contracts.StructuredCompletionRequest {
	return contracts.StructuredCompletionRequest{
		PromptID: "adapter.test", SessionID: "observation-test",
		Inputs: contracts.LLMInputSet{"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{}`), "", "")},
	}
}

func TestPromptKitObservationsComposeTransientRetriesAndRepair(t *testing.T) {
	var calls atomic.Int32
	client := newObservedHTTPClient(t, func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		switch calls.Add(1) {
		case 1, 3:
			w.WriteHeader(http.StatusServiceUnavailable)
			_, _ = io.WriteString(w, `{"error":{"message":"try again"}}`)
		case 2:
			_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"{\"bad\":true}"}}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5,"prompt_tokens_details":{"cached_tokens":0}}}`)
		case 4:
			_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"{\"ok\":true}"}}],"usage":{"prompt_tokens":7,"completion_tokens":0,"total_tokens":7}}`)
		default:
			t.Error("unexpected extra HTTP attempt")
			w.WriteHeader(http.StatusBadRequest)
		}
	}, 1, nil)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	req := observedRequest()
	repairs := 1
	req.StructuredOutputRepairAttempts = &repairs
	var out struct{ OK bool }
	response, err := client.CompleteStructured(ctx, req, &out)
	if err != nil || !out.OK {
		t.Fatalf("completion = (%+v, %v), output = %+v", response, err, out)
	}
	summary := response.Execution
	if calls.Load() != 4 || summary == nil || summary.Attempts != 4 || summary.TransientRetries != 2 || summary.RepairCalls != 1 || summary.Outcome != "success" || summary.RetriesExhausted {
		t.Fatalf("calls = %d, execution = %+v", calls.Load(), summary)
	}
	wantUsage := contracts.LLMObservedTokenUsage{
		PromptTokens:     contracts.LLMObservedTokenCount{Value: 10, Known: true},
		CompletionTokens: contracts.LLMObservedTokenCount{Value: 2, Known: true},
		TotalTokens:      contracts.LLMObservedTokenCount{Value: 12, Known: true},
		CachedTokens:     contracts.LLMObservedTokenCount{Known: true},
	}
	// Failed HTTP attempts contribute unknown usage, not known zero. Neither
	// terminal totals nor legacy result totals should be added a second time.
	if summary.Usage != wantUsage {
		t.Fatalf("observed usage = %+v, want %+v", summary.Usage, wantUsage)
	}
	if response.PromptTokens != 10 || response.CompletionTokens != 2 || response.TotalTokens != 12 || response.RepairAttempts != 1 {
		t.Fatalf("legacy result accounting changed: %+v", response)
	}
}

type closeSignalingBody struct {
	io.ReadCloser
	closed chan<- struct{}
}

func (b closeSignalingBody) Close() error {
	err := b.ReadCloser.Close()
	b.closed <- struct{}{}
	return err
}

func TestPromptKitObservationSurvivesCanceledRetryWait(t *testing.T) {
	consumed := make(chan struct{}, 2)
	transport := http.DefaultTransport.(*http.Transport).Clone()
	t.Cleanup(transport.CloseIdleConnections)
	var calls atomic.Int32
	client := newObservedHTTPClient(t, func(w http.ResponseWriter, r *http.Request) {
		calls.Add(1)
		w.Header().Set("Retry-After", "3600")
		w.WriteHeader(http.StatusServiceUnavailable)
		_, _ = io.WriteString(w, `{"error":{"message":"busy"}}`)
	}, 1, roundTripFunc(func(req *http.Request) (*http.Response, error) {
		response, err := transport.RoundTrip(req)
		if err == nil {
			response.Body = closeSignalingBody{ReadCloser: response.Body, closed: consumed}
		}
		return response, err
	}))
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	type result struct {
		response contracts.StructuredCompletionResponse
		err      error
	}
	done := make(chan result, 1)
	go func() {
		var out map[string]any
		response, err := client.CompleteStructured(ctx, observedRequest(), &out)
		done <- result{response, err}
	}()
	// The failure has been consumed, so cancellation must interrupt the pending
	// retry wait rather than depend on another provider response or its timeout.
	select {
	case <-consumed:
	case <-time.After(5 * time.Second):
		t.Fatal("provider response was not consumed")
	}
	cancel()
	select {
	case got := <-done:
		if !errors.Is(got.err, context.Canceled) || got.response.Execution == nil {
			t.Fatalf("canceled completion = (%+v, %v)", got.response, got.err)
		}
		summary := got.response.Execution
		if summary.Outcome != "error" || summary.FailureKind != "canceled" || summary.Attempts != 1 || summary.RetriesExhausted || calls.Load() != 1 {
			t.Fatalf("canceled execution = %+v, HTTP calls = %d", summary, calls.Load())
		}
	case <-time.After(5 * time.Second):
		t.Fatal("cancellation did not interrupt retry waiting")
	}
}

func TestPromptKitObservationsSeparatePreparationAndExecutionFailure(t *testing.T) {
	var calls atomic.Int32
	client := newObservedHTTPClient(t, func(w http.ResponseWriter, r *http.Request) {
		calls.Add(1)
		w.WriteHeader(http.StatusServiceUnavailable)
		_, _ = io.WriteString(w, `{"error":{"message":"busy"}}`)
	}, 1, nil)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	req := observedRequest()
	req.PromptID = "missing-prompt"
	var out map[string]any
	response, err := client.CompleteStructured(ctx, req, &out)
	if err == nil || response.Execution != nil || calls.Load() != 0 {
		t.Fatalf("preparation failure = (%+v, %v), calls = %d", response, err, calls.Load())
	}
	response, err = client.CompleteStructured(ctx, observedRequest(), &out)
	var generation *contracts.LLMGenerationError
	if !errors.As(err, &generation) || response.Execution == nil {
		t.Fatalf("execution failure = (%+v, %v)", response, err)
	}
	summary := response.Execution
	if summary.Outcome != "error" || summary.FailureKind != "provider_status" || summary.StatusCode != http.StatusServiceUnavailable || !summary.RetriesExhausted || summary.Attempts != 2 || summary.TransientRetries != 1 || calls.Load() != 2 {
		t.Fatalf("failed execution = %+v, calls = %d", summary, calls.Load())
	}
	if summary.Usage != (contracts.LLMObservedTokenUsage{}) {
		t.Fatalf("failed attempts should have unknown usage: %+v", summary.Usage)
	}
}

func TestPromptKitObservationsAreIsolatedAcrossConcurrentCalls(t *testing.T) {
	arrived := make(chan struct{}, 2)
	release := make(chan struct{})
	client := newObservedHTTPClient(t, func(w http.ResponseWriter, r *http.Request) {
		body, err := io.ReadAll(r.Body)
		if err != nil {
			t.Error(err)
			return
		}
		tokens := 11
		if strings.Contains(string(body), "second-call") {
			tokens = 22
		}
		arrived <- struct{}{}
		select {
		case <-release:
		case <-r.Context().Done():
			return
		}
		_ = json.NewEncoder(w).Encode(map[string]any{
			"choices": []any{map[string]any{"message": map[string]string{"content": `{"ok":true}`}}},
			"usage":   map[string]int{"prompt_tokens": tokens, "completion_tokens": 0, "total_tokens": tokens},
		})
	}, 0, nil)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	type result struct {
		want     int
		response contracts.StructuredCompletionResponse
		err      error
	}
	done := make(chan result, 2)
	for name, tokens := range map[string]int{"first-call": 11, "second-call": 22} {
		go func() {
			req := observedRequest()
			req.Vars = map[string]any{"custom": name}
			var out map[string]any
			response, err := client.CompleteStructured(ctx, req, &out)
			done <- result{tokens, response, err}
		}()
	}
	for i := 0; i < 2; i++ {
		select {
		case <-arrived:
		case <-ctx.Done():
			t.Fatal("requests did not overlap")
		}
	}
	close(release)
	for i := 0; i < 2; i++ {
		select {
		case got := <-done:
			summary := got.response.Execution
			if got.err != nil || summary == nil {
				t.Fatalf("concurrent completion = (%+v, %v)", got.response, got.err)
			}
			if summary.Attempts != 1 || summary.Usage.PromptTokens != (contracts.LLMObservedTokenCount{Value: got.want, Known: true, Complete: true}) || summary.Usage.CompletionTokens != (contracts.LLMObservedTokenCount{Known: true, Complete: true}) || summary.Usage.CachedTokens != (contracts.LLMObservedTokenCount{}) {
				t.Fatalf("execution for %d tokens = %+v", got.want, summary)
			}
		case <-ctx.Done():
			t.Fatal("concurrent completions did not return")
		}
	}
}