package main

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"strings"
	"sync"
	"testing"

	ck "gitea.maximumdirect.net/eric/ConfigKit"
	pk "gitea.maximumdirect.net/eric/promptkit"
)

func TestDemonstrationSelectsOutboundModels(t *testing.T) {
	for _, tc := range []struct {
		name                          string
		args                          []string
		overlay, proposal, assessment string
	}{
		{"base", nil, "", "synthetic-grammar", "synthetic-review"},
		{"testing", []string{"--overlay", "testing"}, "testing", "synthetic-grammar-testing", "synthetic-review-testing"},
		{"proposal override", []string{"--overlay", "testing", "--proposal-profile", "proposal"}, "testing", "synthetic-grammar", "synthetic-review-testing"},
		{"assessment override", []string{"--overlay", "testing", "--validation-profile", "assessment"}, "testing", "synthetic-grammar-testing", "synthetic-review"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			t.Setenv("AUDITA_PROTOTYPE_KEY", "synthetic-test-secret")
			var mu sync.Mutex
			var models []string
			base := demonstrationClient()
			client := &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) {
				body, err := io.ReadAll(req.Body)
				if err != nil {
					return nil, err
				}
				var wire struct {
					Model string `json:"model"`
				}
				if err := json.Unmarshal(body, &wire); err != nil {
					return nil, err
				}
				mu.Lock()
				models = append(models, wire.Model)
				mu.Unlock()
				req.Body = io.NopCloser(bytes.NewReader(body))
				return base.Transport.RoundTrip(req)
			})}
			var out bytes.Buffer
			if err := run(tc.args, &out, client); err != nil {
				t.Fatal(err)
			}
			var got struct {
				Overlay    string                                   `json:"overlay"`
				Proposal   struct{ Profile, Backend, Model string } `json:"proposal"`
				Assessment struct{ Profile, Backend, Model string } `json:"assessment"`
				Result     WorkflowResult                           `json:"result"`
				Attempts   uint64                                   `json:"attempts"`
				Retries    uint64                                   `json:"transient_retries"`
				Repairs    uint64                                   `json:"structural_repairs"`
			}
			if err := json.Unmarshal(out.Bytes(), &got); err != nil {
				t.Fatal(err)
			}
			if got.Overlay != tc.overlay || got.Proposal.Model != tc.proposal || got.Assessment.Model != tc.assessment || got.Proposal.Backend != "offline" || got.Assessment.Backend != "offline" {
				t.Fatalf("selection = %+v", got)
			}
			profiles := map[string]string{"synthetic-grammar": "proposal", "synthetic-review": "assessment", "synthetic-grammar-testing": "proposal-testing", "synthetic-review-testing": "assessment-testing"}
			if got.Proposal.Profile != profiles[tc.proposal] || got.Assessment.Profile != profiles[tc.assessment] {
				t.Fatalf("reported profiles = %q, %q", got.Proposal.Profile, got.Assessment.Profile)
			}
			if got.Result.Accepted != 1 || got.Attempts != 4 || got.Retries != 0 || got.Repairs != 2 {
				t.Fatalf("outcome = %+v", got)
			}
			mu.Lock()
			defer mu.Unlock()
			if len(models) != 4 || models[0] != tc.proposal || models[1] != tc.proposal || models[2] != tc.assessment || models[3] != tc.assessment {
				t.Fatalf("outbound models = %v", models)
			}
		})
	}
}

type step struct {
	status  int
	content string
	usage   any
}

func testRuntime(t *testing.T, transport http.RoundTripper, overrides ...ck.Override) *Runtime {
	t.Helper()
	t.Setenv("AUDITA_PROTOTYPE_KEY", "synthetic-test-secret")
	loaded, err := loadConfig(ck.File("config.yaml"), "", overrides...)
	if err != nil {
		t.Fatal(err)
	}
	r, err := newRuntime(loaded.Config, ".", &http.Client{Transport: transport})
	if err != nil {
		t.Fatal(err)
	}
	return r
}
func intOverride(field string, value int) ck.Override {
	return ck.Override{Path: ck.Keys("llm", field), Value: value, Flag: "--" + field}
}

func TestWorkflowBoundaries(t *testing.T) {
	zero := map[string]any{"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
	for _, tc := range []struct {
		name                              string
		steps                             []step
		retries, repairs                  int
		accepted, rejected, warnings      int
		attempts, retryCount, repairCount uint64
		cause                             error
	}{
		{name: "repair both roles", steps: []step{{200, "bad", zero}, {200, validProposal, zero}, {200, "bad", nil}, {200, validAssessment, zero}}, repairs: 2, accepted: 1, attempts: 4, repairCount: 2},
		{name: "proposal repair exhausted", steps: []step{{200, "bad", nil}, {200, "bad", nil}, {200, "bad", nil}}, repairs: 2, warnings: 1, attempts: 3, repairCount: 2},
		{name: "assessment repair exhausted", steps: []step{{200, validProposal, nil}, {200, "bad", nil}, {200, "bad", nil}}, repairs: 1, rejected: 1, warnings: 1, attempts: 3, repairCount: 1},
		{name: "domain rejects missing span", steps: []step{{200, strings.Replace(validProposal, "hello world", "absent span", 1), nil}}, repairs: 2, rejected: 1, attempts: 1},
		{name: "domain rejects missing decision", steps: []step{{200, validProposal, nil}, {200, `{"validations":[]}`, nil}}, repairs: 2, rejected: 1, warnings: 1, attempts: 2},
		{name: "domain rejects duplicate decision", steps: []step{{200, validProposal, nil}, {200, `{"validations":[{"correction_index":0,"approved":true,"confidence":1,"reason":"ok"},{"correction_index":0,"approved":true,"confidence":1,"reason":"ok"}]}`, nil}}, repairs: 2, rejected: 1, warnings: 1, attempts: 2},
		{name: "assessment semantic rejection", steps: []step{{200, validProposal, nil}, {200, strings.Replace(validAssessment, `"approved":true`, `"approved":false`, 1), nil}}, repairs: 2, rejected: 1, attempts: 2},
		{name: "429 recovery", steps: []step{{429, "", nil}, {200, validProposal, zero}, {200, validAssessment, zero}}, retries: 1, accepted: 1, attempts: 3, retryCount: 1},
		{name: "explicit zero retries", steps: []step{{503, "", nil}}, repairs: 2, attempts: 1, cause: pk.ErrLLMGenerate},
		{name: "retry exhaustion after completed proposal", steps: []step{{200, validProposal, zero}, {503, "", nil}, {503, "", nil}}, retries: 1, repairs: 2, attempts: 3, retryCount: 1, cause: pk.ErrRetriesExhausted},
		{name: "retry and repair budgets", steps: []step{{503, "", nil}, {200, "bad", zero}, {503, "", nil}, {200, validProposal, zero}, {200, validAssessment, zero}}, retries: 1, repairs: 1, accepted: 1, attempts: 5, retryCount: 2, repairCount: 1},
	} {
		t.Run(tc.name, func(t *testing.T) {
			var mu sync.Mutex
			count := 0
			transport := roundTripFunc(func(req *http.Request) (*http.Response, error) {
				mu.Lock()
				defer mu.Unlock()
				if req.Header.Get("Authorization") != "Bearer synthetic-test-secret" {
					t.Error("wrong credential")
				}
				var body struct {
					Messages []struct {
						Content string `json:"content"`
					} `json:"messages"`
				}
				if err := json.NewDecoder(req.Body).Decode(&body); err != nil {
					return nil, err
				}
				if len(body.Messages) < 2 || !strings.Contains(body.Messages[1].Content, "Prompt hardening policy") {
					t.Error("real prompt hardening missing")
				}
				if count >= len(tc.steps) {
					t.Error("unexpected extra provider call")
					return providerResponse(400, "", nil), nil
				}
				s := tc.steps[count]
				count++
				return providerResponse(s.status, s.content, s.usage), nil
			})
			r := testRuntime(t, transport, intOverride("transient_retries", tc.retries), intOverride("repair_attempts", tc.repairs))
			recorder := &Recorder{}
			out, err := r.workflow(context.Background(), recorder.observer)
			if tc.cause != nil {
				if !errors.Is(err, tc.cause) {
					t.Fatalf("error %v does not match %v", err, tc.cause)
				}
			} else if err != nil {
				t.Fatal(err)
			}
			if out.Accepted != tc.accepted || out.Rejected != tc.rejected || len(out.Warnings) != tc.warnings {
				t.Fatalf("result %+v", out)
			}
			a, rt, rp := recorder.totals()
			if a != tc.attempts || rt != tc.retryCount || rp != tc.repairCount || count != len(tc.steps) {
				t.Fatalf("counts %d/%d/%d; dispatches %d", a, rt, rp, count)
			}
			records := recorder.snapshot()
			serialized, _ := json.Marshal(records)
			if strings.Contains(string(serialized), "synthetic-test-secret") || strings.Contains(string(serialized), "hello world") {
				t.Fatal("sensitive observation")
			}
			for _, o := range records {
				if o.Kind == string(pk.ExecutionAttempt) && o.StatusCode == 503 && o.Usage.Total.Known {
					t.Fatal("invented known usage")
				}
			}
		})
	}
}