package llm

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

	"gitea.maximumdirect.net/eric/audita/v3/internal/core/config"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/responseschema"
	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/structuredoutput"
	"gitea.maximumdirect.net/eric/audita/v3/internal/prompts"
	ck "go.maximumdirect.net/configkit"
	pk "go.maximumdirect.net/promptkit"
)

func TestNativeConversationReplayWithRepairContract(t *testing.T) {
	dir := t.TempDir()
	files := map[string]string{
		"config.yml": "version: 3\n",
		"prompts.yml": `profiles:
  - id: proposal-profile
    endpoint: https://example.invalid/v1
    model: proposal-model
  - id: repair-profile
    endpoint: https://example.invalid/v1
    model: repair-model
prompts:
  - id: proposal-fixture
    messages:
      - role: system
        content: 'System {{.Subject}}'
      - role: user
        content: 'Proposal {{.Subject}}'
    output:
      format: json
      validation_mode: json_schema
      schema_path: correction_set.v1.json
  - id: repair-fixture
    messages:
      - role: user
        content: 'Repair {{.Finding}}'
    output:
      format: text
      validation_mode: none
`,
	}
	for name, content := range files {
		if err := os.WriteFile(filepath.Join(dir, name), []byte(content), 0644); err != nil {
			t.Fatal(err)
		}
	}
	loaded, err := config.Load(ck.Input{Source: ck.File(filepath.Join(dir, "config.yml"))}, "")
	if err != nil {
		t.Fatal(err)
	}
	var requests []struct {
		Model          string                           `json:"model"`
		Messages       []struct{ Role, Content string } `json:"messages"`
		ResponseFormat map[string]any                   `json:"response_format"`
	}
	outputs := []string{`{"corrections":[]}`, `{"proposal_ref":1,"action":"revise","correction":{"id":1,"original_text":" ","corrected_text":"","confidence":0.8}}`, `{"proposal_ref":1,"action":"withdraw","correction":null}`}
	transport := transportFunc(func(r *http.Request) (*http.Response, error) {
		var request struct {
			Model          string                           `json:"model"`
			Messages       []struct{ Role, Content string } `json:"messages"`
			ResponseFormat map[string]any                   `json:"response_format"`
		}
		if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
			t.Fatal(err)
		}
		requests = append(requests, request)
		content, _ := json.Marshal(outputs[len(requests)-1])
		body := fmt.Sprintf(`{"model":%q,"choices":[{"message":{"content":%s}}]}`, request.Model, content)
		return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body)), Request: r}, nil
	})
	runtime, err := NewRuntime(context.Background(), loaded, &http.Client{Transport: transport})
	if err != nil {
		t.Fatal(err)
	}
	ctx := context.Background()
	base := pk.RunRequest{PromptID: "proposal-fixture", ProfileID: "proposal-profile", Vars: map[string]string{"Subject": "source"}}
	original, err := runtime.Engine.PrepareExecution(ctx, base)
	if err != nil {
		t.Fatal(err)
	}
	prefix := original.Details().Messages
	initial, err := runtime.Engine.RunPrepared(ctx, original)
	if err != nil {
		t.Fatal(err)
	}
	suffix := []pk.RenderedMessage{{Role: pk.RoleAssistant, Content: initial.RawOutput}}
	contract, err := OutputContract(responseschema.CorrectionRepairKey, 0)
	if err != nil {
		t.Fatal(err)
	}
	for _, finding := range []string{"first", "second"} {
		instruction, err := runtime.Engine.Prepare(ctx, pk.RunRequest{PromptID: "repair-fixture", ProfileID: "repair-profile", Vars: map[string]string{"Finding": finding}})
		if err != nil {
			t.Fatal(err)
		}
		suffix = append(suffix, instruction.Messages...)
		req := base
		req.ProfileID = "repair-profile"
		req.AppendedMessages = suffix
		req.Validation = contract
		handle, err := runtime.Engine.PrepareExecution(ctx, req)
		if err != nil {
			t.Fatal(err)
		}
		details := handle.Details()
		if !reflect.DeepEqual(details.Messages[:len(prefix)], prefix) || !reflect.DeepEqual(details.Messages[len(prefix):], suffix) {
			t.Fatalf("prepared conversation = %#v", details.Messages)
		}
		result, err := runtime.Engine.RunPrepared(ctx, handle)
		if err != nil {
			t.Fatal(err)
		}
		if !result.Validation.IsValid {
			t.Fatalf("invalid result: %+v", result.Validation)
		}
		suffix = append(suffix, pk.RenderedMessage{Role: pk.RoleAssistant, Content: result.RawOutput})
	}
	if len(requests) != 3 || requests[0].Model != "proposal-model" || requests[1].Model != "repair-model" || requests[2].Model != "repair-model" {
		t.Fatalf("models: %+v", requests)
	}
	for i, wantLen := range []int{2, 4, 6} {
		req := requests[i]
		if len(req.Messages) != wantLen {
			t.Fatalf("request %d messages: %+v", i, req.Messages)
		}
		for j, message := range req.Messages {
			if message.Role != string(append(prefix, suffix...)[j].Role) || message.Content != append(prefix, suffix...)[j].Content {
				t.Fatalf("request %d message %d: %+v", i, j, message)
			}
		}
		encoded, _ := json.Marshal(req.ResponseFormat)
		expected := "corrections"
		if i > 0 {
			expected = "proposal_ref"
		}
		if !strings.Contains(string(encoded), expected) {
			t.Fatalf("request %d contract: %s", i, encoded)
		}
	}
}

func TestRoleClientCapturesAndReplaysRepairConversation(t *testing.T) {
	dir := t.TempDir()
	for name, content := range map[string]string{
		"config.yml":  "version: 3\n",
		"prompts.yml": "profiles:\n  - id: proposal\n    endpoint: https://example.invalid/v1\n    model: proposal\n  - id: repair\n    endpoint: https://example.invalid/v1\n    model: repair\ncollections:\n  audita:\n    prompts:\n      - id: modules.grammar.repair\n        messages:\n          - role: user\n            content: '{{ .RepairRequestJSON }}'\n            cache_control: {type: ephemeral, ttl: 1h}\n        output:\n          format: json\n          validation_mode: json_schema\n          schema_path: correction_repair.v1.json\n",
	} {
		if err := os.WriteFile(filepath.Join(dir, name), []byte(content), 0644); err != nil {
			t.Fatal(err)
		}
	}
	loaded, err := config.Load(ck.Input{Source: ck.File(filepath.Join(dir, "config.yml"))}, "")
	if err != nil {
		t.Fatal(err)
	}
	type wireMessage struct {
		Role    string
		Content json.RawMessage
	}
	var messages [][]wireMessage
	var models []string
	outputs := []string{`{"corrections":[]}`, `{"proposal_ref":1,"action":"revise","correction":{"id":1,"original_text":"a","corrected_text":"b","confidence":0.8}}`, `{"proposal_ref":1,"action":"withdraw","correction":null}`, `{"proposal_ref":1,"action":"withdraw","correction":null}`}
	transport := transportFunc(func(r *http.Request) (*http.Response, error) {
		var req struct {
			Model    string        `json:"model"`
			Messages []wireMessage `json:"messages"`
		}
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			t.Fatal(err)
		}
		models = append(models, req.Model)
		messages = append(messages, req.Messages)
		content, _ := json.Marshal(outputs[len(messages)-1])
		body := fmt.Sprintf(`{"model":%q,"choices":[{"message":{"content":%s}}]}`, req.Model, content)
		return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body)), Request: r}, nil
	})
	runtime, err := NewRuntime(context.Background(), loaded, &http.Client{Transport: transport})
	if err != nil {
		t.Fatal(err)
	}
	settings := config.LLM{OperationTimeout: 10 * time.Second}
	proposal, err := runtime.NewRoleClient("proposal", "proposal", settings, CallHooks{})
	if err != nil {
		t.Fatal(err)
	}
	repair, err := runtime.NewRoleClient("repair", "repair", settings, CallHooks{})
	if err != nil {
		t.Fatal(err)
	}
	vars := map[string]string{"TranscriptDescriptionBlock": "", "GlossaryJSON": "[]", "SectionJSON": "[]"}
	initial, err := proposal.CompleteStructured(context.Background(), structuredoutput.Request{PromptID: prompts.PromptIDModuleGrammarProposal, Schema: responseschema.CorrectionSetKey, Variables: vars, CaptureExchange: true}, &map[string]any{})
	if err != nil || initial.Exchange == nil {
		t.Fatalf("capture: %+v %v", initial, err)
	}
	vars["SectionJSON"] = "changed"
	replay := &structuredoutput.Replay{Exchange: initial.Exchange}
	for _, finding := range []string{"first", "second"} {
		result, err := repair.CompleteStructured(context.Background(), structuredoutput.Request{PromptID: prompts.PromptIDModuleGrammarRepair, Schema: responseschema.CorrectionRepairKey, Variables: map[string]string{"RepairRequestJSON": finding}, Replay: replay}, &map[string]any{})
		if err != nil {
			t.Fatal(err)
		}
		if len(result.RepairInstruction) != 1 {
			t.Fatalf("instruction: %+v", result)
		}
		replay.Turns = append(replay.Turns, result.RepairInstruction...)
		replay.Turns = append(replay.Turns, structuredoutput.Message{Role: "assistant", Content: result.FinalContent})
	}
	if !reflect.DeepEqual(models, []string{"proposal", "repair", "repair"}) {
		t.Fatalf("models = %v", models)
	}
	prefix := messages[0]
	if len(messages) != 3 || !reflect.DeepEqual(messages[1][:len(prefix)], prefix) || !reflect.DeepEqual(messages[2][:len(prefix)], prefix) {
		t.Fatalf("prefixes: %+v", messages)
	}
	encodedRevision, _ := json.Marshal(outputs[1])
	if len(messages[1]) != len(prefix)+2 || len(messages[2]) != len(prefix)+4 || string(messages[2][len(prefix)+2].Content) != string(encodedRevision) {
		t.Fatalf("turns: %+v", messages)
	}
	first := messages[1][len(prefix)+1]
	if first.Role != "user" || !strings.Contains(string(first.Content), `"ttl":"1h"`) ||
		!strings.Contains(string(first.Content), `"text":"first"`) || !reflect.DeepEqual(first, messages[2][len(prefix)+1]) {
		t.Fatalf("cached instruction changed during replay: first=%s replay=%s", first.Content, messages[2][len(prefix)+1].Content)
	}
	_, err = repair.CompleteStructured(context.Background(), structuredoutput.Request{PromptID: prompts.PromptIDModuleGrammarRepair, Schema: responseschema.CorrectionRepairKey, Variables: map[string]string{"RepairRequestJSON": "independent"}, Replay: &structuredoutput.Replay{Exchange: initial.Exchange}}, &map[string]any{})
	if err != nil || len(messages) != 4 || len(messages[3]) != len(prefix)+2 || !strings.Contains(string(messages[3][len(prefix)+1].Content), "independent") {
		t.Fatalf("independent branch: %v %+v", err, messages)
	}
	bad := *initial.Exchange
	bad.Messages = append([]structuredoutput.Message(nil), bad.Messages...)
	bad.Messages[0].Content = "changed"
	_, err = repair.CompleteStructured(context.Background(), structuredoutput.Request{PromptID: prompts.PromptIDModuleGrammarRepair, Schema: responseschema.CorrectionRepairKey, Variables: map[string]string{"RepairRequestJSON": "failure"}, Replay: &structuredoutput.Replay{Exchange: &bad}}, &map[string]any{})
	var failure *structuredoutput.Failure
	if !errors.As(err, &failure) || failure.Class != structuredoutput.FailureFatal || len(messages) != 4 {
		t.Fatalf("prefix failure: %v, calls %d", err, len(messages))
	}
}

func TestRepairReplayUsesReplacementAndCountsStructuralRepair(t *testing.T) {
	dir := t.TempDir()
	files := map[string]string{
		"config.yml":  "version: 3\n",
		"prompts.yml": "profiles:\n  - id: fixture\n    endpoint: https://example.invalid/v1\n    model: fixture\ncollections:\n  audita:\n    prompts:\n      - id: modules.grammar.repair\n        messages:\n          - role: user\n            content_file: repair.txt\n        output:\n          format: json\n          validation_mode: json_schema\n          schema_path: correction_repair.v1.json\n",
		"repair.txt":  "Replace {{.RepairRequestJSON}}",
	}
	for name, content := range files {
		if err := os.WriteFile(filepath.Join(dir, name), []byte(content), 0644); err != nil {
			t.Fatal(err)
		}
	}
	loaded, err := config.Load(ck.Input{Source: ck.File(filepath.Join(dir, "config.yml"))}, "")
	if err != nil {
		t.Fatal(err)
	}
	outputs := []string{`{"corrections":[]}`, `{"bad":true}`, `{"proposal_ref":1,"action":"withdraw","correction":null}`}
	var requests [][]structuredoutput.Message
	transport := transportFunc(func(r *http.Request) (*http.Response, error) {
		var req struct {
			Messages []structuredoutput.Message `json:"messages"`
		}
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			t.Fatal(err)
		}
		requests = append(requests, req.Messages)
		content, _ := json.Marshal(outputs[len(requests)-1])
		body := fmt.Sprintf(`{"model":"fixture","choices":[{"message":{"content":%s}}]}`, content)
		return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body)), Request: r}, nil
	})
	runtime, err := NewRuntime(context.Background(), loaded, &http.Client{Transport: transport})
	if err != nil {
		t.Fatal(err)
	}
	client, err := runtime.NewRoleClient("repair", "fixture", config.LLM{OperationTimeout: 10 * time.Second, RepairAttempts: 1}, CallHooks{})
	if err != nil {
		t.Fatal(err)
	}
	base, err := client.CompleteStructured(context.Background(), structuredoutput.Request{PromptID: prompts.PromptIDModuleGrammarProposal, Schema: responseschema.CorrectionSetKey, Variables: map[string]string{"TranscriptDescriptionBlock": "", "GlossaryJSON": "[]", "SectionJSON": "[]"}, CaptureExchange: true}, &map[string]any{})
	if err != nil {
		t.Fatal(err)
	}
	result, err := client.CompleteStructured(context.Background(), structuredoutput.Request{PromptID: prompts.PromptIDModuleGrammarRepair, Schema: responseschema.CorrectionRepairKey, Variables: map[string]string{"RepairRequestJSON": "fixture"}, Replay: &structuredoutput.Replay{Exchange: base.Exchange}}, &map[string]any{})
	if err != nil {
		t.Fatal(err)
	}
	if len(requests) != 3 || result.Execution.ProviderAttempts != 2 || result.Execution.StructuralRepairs != 1 || result.FinalContent != outputs[2] || len(result.RepairInstruction) != 1 || result.RepairInstruction[0].Content != "Replace fixture" {
		t.Fatalf("result=%+v calls=%d", result, len(requests))
	}
}