package main

import (
	"context"
	"embed"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"time"

	"go.maximumdirect.net/configkit"
	"go.maximumdirect.net/promptkit"
)

//go:embed prompts/*.yaml schemas/*.json
var assets embed.FS

type scriptedTransport struct {
	proposalCalls int
}

func (t *scriptedTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	if err := req.Context().Err(); err != nil {
		return nil, err
	}
	var request struct {
		Model string `json:"model"`
	}
	if err := json.NewDecoder(req.Body).Decode(&request); err != nil {
		return nil, err
	}
	switch request.Model {
	case "proposal-model":
		t.proposalCalls++
		if t.proposalCalls == 1 {
			return httpResponse(http.StatusServiceUnavailable, `{"error":{}}`), nil
		}
		return modelResponse(`{"title":"Queue cleanup","risk":"low"}`, 10, 8), nil
	case "assessment-model":
		return modelResponse(`{"recommendation":"accept","confidence":0.9}`, 7, 5), nil
	default:
		return nil, fmt.Errorf("unexpected model in offline transport")
	}
}

func modelResponse(content string, promptTokens, completionTokens int) *http.Response {
	body, _ := json.Marshal(map[string]any{
		"choices": []any{map[string]any{"message": map[string]string{"content": content}}},
		"usage":   map[string]int{"prompt_tokens": promptTokens, "completion_tokens": completionTokens},
	})
	return httpResponse(http.StatusOK, string(body))
}

func httpResponse(status int, body string) *http.Response {
	return &http.Response{StatusCode: status, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body))}
}

type eventLog struct{ events []promptkit.ExecutionEvent }

func (l *eventLog) record(event promptkit.ExecutionEvent) { l.events = append(l.events, event) }

func (l *eventLog) terminal() (promptkit.ExecutionEvent, error) {
	if len(l.events) == 0 || l.events[len(l.events)-1].Kind != promptkit.ExecutionFinished {
		return promptkit.ExecutionEvent{}, fmt.Errorf("execution did not finish")
	}
	return l.events[len(l.events)-1], nil
}

type summary struct {
	ProposalAttempts           uint64                        `json:"proposal_attempts"`
	ProposalTransientRetries   uint64                        `json:"proposal_transient_retries"`
	ProposalUsageComplete      bool                          `json:"proposal_usage_complete"`
	AssessmentAttempts         uint64                        `json:"assessment_attempts"`
	AssessmentTransientRetries uint64                        `json:"assessment_transient_retries"`
	AssessmentOutcome          string                        `json:"assessment_outcome"`
	CoordinationOutcome        promptkit.CoordinationOutcome `json:"coordination_outcome"`
	ExhaustedOutcome           promptkit.CoordinationOutcome `json:"exhausted_outcome"`
	GroupMembers               int                           `json:"group_members"`
	GroupAttempts              uint64                        `json:"group_attempts"`
	PreparedMessages           int                           `json:"prepared_messages"`
	PreparedParentMatches      bool                          `json:"prepared_parent_matches"`
}

func run() (summary, error) {
	transport := &scriptedTransport{}
	inputs, err := promptkit.FilesystemInputs("example prompts", assets, "prompts")
	if err != nil {
		return summary{}, err
	}
	backendID, proposalModel, assessmentModel := "offline-backend", "proposal-model", "assessment-model"
	backend := promptkit.LocalBackend("https://example.invalid/v1", 1)
	backend.ID = backendID
	inputs = append(inputs,
		configkit.Input{Source: promptkit.BackendSource("example backend", []promptkit.BackendConfig{
			backend,
		})},
		configkit.Input{Source: promptkit.ProfileSource("example profiles", []promptkit.ProfileConfig{
			{ID: "proposal-profile", Backend: &backendID, Model: &proposalModel},
			{ID: "assessment-profile", Backend: &backendID, Model: &assessmentModel},
		})},
	)
	engine, err := promptkit.NewEngine(promptkit.Config{
		Sources:    promptkit.ConfigurationSources{Underlays: []configkit.Underlay{{Inputs: inputs}}},
		HTTPClient: &http.Client{Transport: transport},
	}, promptkit.WithBundledOpenRouter(false), promptkit.WithSchemaFS(assets, "schemas")) // Keep the example's configuration self-contained.
	if err != nil {
		return summary{}, err
	}

	var proposalLog eventLog
	var executions promptkit.ExecutionGroup
	proposal, err := engine.Run(context.Background(), promptkit.RunRequest{
		PromptID: "example.proposal", Inputs: map[string]promptkit.ArtifactRef{"idea": promptkit.Inline("Reduce queue cleanup toil.")}, CaptureExchange: true, Observer: promptkit.ComposeObservers(proposalLog.record, executions.Observe),
	})
	if err != nil {
		return summary{}, err
	}
	if !proposal.Validation.IsValid {
		return summary{}, fmt.Errorf("proposal failed structural validation")
	}
	proposalTerminal, err := proposalLog.terminal()
	if err != nil {
		return summary{}, err
	}
	var assessmentLog eventLog
	var preparedDetails promptkit.PreparedRun
	decision := "rejected"
	coordination, err := engine.Coordinate(context.Background(), promptkit.CoordinateRequest{
		Start: proposal.Exchange, AdditionalTurns: 1, Observer: executions.Observe,
		Next: func(ctx context.Context, turn promptkit.ConversationTurn) (*promptkit.RunRequest, error) {
			if turn.Result == nil {
				return &promptkit.RunRequest{PromptID: "example.proposal", Inputs: map[string]promptkit.ArtifactRef{"idea": promptkit.Inline("Refine the proposal.")}, OperationTimeout: 10 * time.Second, OnPrepared: func(_ context.Context, details promptkit.PreparedRun) { preparedDetails = details }}, nil
			}
			if !turn.Result.Validation.IsValid {
				return nil, fmt.Errorf("follow-up failed structural validation")
			}
			assessment, err := engine.Run(ctx, promptkit.RunRequest{PromptID: "example.assessment", Inputs: map[string]promptkit.ArtifactRef{"proposal": promptkit.Inline(turn.Result.RawOutput)}, OperationTimeout: 10 * time.Second, Observer: promptkit.ComposeObservers(assessmentLog.record, executions.Observe)})
			if err != nil {
				return nil, err
			}
			if !assessment.Validation.IsValid {
				return nil, fmt.Errorf("assessment failed structural validation")
			}
			var proposalValue struct {
				Risk string `json:"risk"`
			}
			var assessmentValue struct {
				Recommendation string `json:"recommendation"`
			}
			if err := json.Unmarshal([]byte(turn.Result.RawOutput), &proposalValue); err != nil {
				return nil, err
			}
			if err := json.Unmarshal([]byte(assessment.RawOutput), &assessmentValue); err != nil {
				return nil, err
			}
			if proposalValue.Risk == "low" && assessmentValue.Recommendation == "accept" {
				decision = "accepted"
			}
			return nil, nil
		},
	})
	if err != nil {
		return summary{}, err
	}
	exhausted, err := engine.Coordinate(context.Background(), promptkit.CoordinateRequest{Start: proposal.Exchange, Next: func(context.Context, promptkit.ConversationTurn) (*promptkit.RunRequest, error) {
		return &promptkit.RunRequest{PromptID: "example.proposal"}, nil
	}})
	if err != nil {
		return summary{}, err
	}
	assessmentTerminal, err := assessmentLog.terminal()
	if err != nil {
		return summary{}, err
	}
	group := executions.Snapshot()
	return summary{
		ProposalAttempts:           proposalTerminal.Attempts,
		ProposalTransientRetries:   proposalTerminal.TransientRetries,
		ProposalUsageComplete:      proposalTerminal.Usage.PromptTokens.Complete,
		AssessmentAttempts:         assessmentTerminal.Attempts,
		AssessmentTransientRetries: assessmentTerminal.TransientRetries,
		AssessmentOutcome:          decision,
		CoordinationOutcome:        coordination.Outcome,
		ExhaustedOutcome:           exhausted.Outcome,
		GroupMembers:               len(group.Members),
		GroupAttempts:              group.Attempts,
		PreparedMessages:           len(preparedDetails.Messages),
		PreparedParentMatches:      preparedDetails.ParentRunID == proposal.RunID,
	}, nil
}

func main() {
	if err := execute(os.Stdout); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func execute(out io.Writer) error {
	value, err := run()
	if err != nil {
		return err
	}
	return json.NewEncoder(out).Encode(value)
}