// Package format formats already-prepared public data for adapters.
package format

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"strings"

	"gitea.maximumdirect.net/eric/promptkit"
)

var ErrUnknownPreparedRunFormat = errors.New("unknown prepared run format")

// PreparedRunOutputFormat is the output format for prepared render data.
type PreparedRunOutputFormat string

const (
	PreparedRunFormatText PreparedRunOutputFormat = "text"
	PreparedRunFormatJSON PreparedRunOutputFormat = "json"

	DefaultPreparedRunOutputFormat PreparedRunOutputFormat = PreparedRunFormatText
)

// PreparedRunFormatter serializes a prepared run without performing use case work.
type PreparedRunFormatter interface {
	Format(prepared *promptkit.PreparedRun) ([]byte, error)
}

// ParsePreparedRunOutputFormat parses a format name.
func ParsePreparedRunOutputFormat(raw string) (PreparedRunOutputFormat, error) {
	switch PreparedRunOutputFormat(strings.ToLower(strings.TrimSpace(raw))) {
	case "":
		return DefaultPreparedRunOutputFormat, nil
	case PreparedRunFormatText:
		return PreparedRunFormatText, nil
	case PreparedRunFormatJSON:
		return PreparedRunFormatJSON, nil
	default:
		return "", fmt.Errorf("%w: %q (supported: %s, %s)", ErrUnknownPreparedRunFormat, raw, PreparedRunFormatText, PreparedRunFormatJSON)
	}
}

// FormatterForPreparedRun returns a formatter strategy for the chosen format.
func FormatterForPreparedRun(outputFormat PreparedRunOutputFormat) (PreparedRunFormatter, error) {
	switch outputFormat {
	case PreparedRunFormatText:
		return textPreparedRunFormatter{}, nil
	case PreparedRunFormatJSON:
		return jsonPreparedRunFormatter{}, nil
	default:
		return nil, fmt.Errorf("%w: %q", ErrUnknownPreparedRunFormat, outputFormat)
	}
}

// FormatPreparedRun formats a prepared run using the selected format.
func FormatPreparedRun(prepared *promptkit.PreparedRun, outputFormat PreparedRunOutputFormat) ([]byte, error) {
	formatter, err := FormatterForPreparedRun(outputFormat)
	if err != nil {
		return nil, err
	}
	return formatter.Format(prepared)
}

// FormatPreparedRunByName parses a format name and formats a prepared run.
func FormatPreparedRunByName(prepared *promptkit.PreparedRun, rawFormat string) ([]byte, error) {
	outputFormat, err := ParsePreparedRunOutputFormat(rawFormat)
	if err != nil {
		return nil, err
	}
	return FormatPreparedRun(prepared, outputFormat)
}

type jsonPreparedRunFormatter struct{}

func (jsonPreparedRunFormatter) Format(prepared *promptkit.PreparedRun) ([]byte, error) {
	if prepared == nil {
		return nil, errors.New("prepared run is nil")
	}
	return json.MarshalIndent(prepared, "", "  ")
}

type textPreparedRunFormatter struct{}

func (textPreparedRunFormatter) Format(prepared *promptkit.PreparedRun) ([]byte, error) {
	if prepared == nil {
		return nil, errors.New("prepared run is nil")
	}

	var b bytes.Buffer
	fmt.Fprintf(&b, "prompt: %s\n", prepared.PromptID)
	fmt.Fprintf(&b, "prompt_version: %s\n", prepared.PromptVersion)
	fmt.Fprintf(&b, "selected_profile_id: %s\n", prepared.SelectedProfileID)
	if prepared.PromptHash != "" {
		fmt.Fprintf(&b, "prompt_hash: %s\n", prepared.PromptHash)
	}
	if prepared.SessionID != "" {
		fmt.Fprintf(&b, "session_id: %s\n", prepared.SessionID)
	}
	fmt.Fprintf(&b, "rendered_prompt_hash: %s\n", prepared.RenderedPromptHash)

	target := prepared.EffectiveModelParams
	fmt.Fprintln(&b, "effective_model_params:")
	fmt.Fprintf(&b, "  endpoint: %s\n", target.Endpoint)
	fmt.Fprintf(&b, "  model: %s\n", target.Model)
	fmt.Fprintf(&b, "  temperature: %g\n", target.Temperature)
	fmt.Fprintf(&b, "  max_tokens: %d\n", target.MaxTokens)
	fmt.Fprintf(&b, "  top_p: %g\n", target.TopP)
	fmt.Fprintf(&b, "  timeout_seconds: %d\n", target.TimeoutSeconds)
	if target.ServiceTier != "" {
		fmt.Fprintf(&b, "  service_tier: %s\n", target.ServiceTier)
	}
	if target.ReasoningEffort != "" {
		fmt.Fprintf(&b, "  reasoning_effort: %s\n", target.ReasoningEffort)
	}
	if target.APIKeyEnv != "" {
		fmt.Fprintf(&b, "  api_key_env: %s\n", target.APIKeyEnv)
	}
	if len(target.ExtraParams) > 0 {
		fmt.Fprintln(&b, "  extra_params:")
		keys := make([]string, 0, len(target.ExtraParams))
		for k := range target.ExtraParams {
			keys = append(keys, k)
		}
		sort.Strings(keys)
		for _, k := range keys {
			renderedValue, err := formatExtraParamTextValue(target.ExtraParams[k])
			if err != nil {
				return nil, fmt.Errorf("failed to format extra_params.%s: %w", k, err)
			}
			fmt.Fprintf(&b, "    %s: %s\n", k, renderedValue)
		}
	}

	fmt.Fprintln(&b, "input_hashes:")
	inputKeys := make([]string, 0, len(prepared.InputHashes))
	for k := range prepared.InputHashes {
		inputKeys = append(inputKeys, k)
	}
	sort.Strings(inputKeys)
	for _, k := range inputKeys {
		fmt.Fprintf(&b, "  %s: %s\n", k, prepared.InputHashes[k])
	}

	fmt.Fprintln(&b, "messages:")
	roleOrder := make([]string, 0)
	byRole := make(map[string][]promptkit.RenderedMessage)
	for _, msg := range prepared.Messages {
		if _, exists := byRole[msg.Role]; !exists {
			roleOrder = append(roleOrder, msg.Role)
		}
		byRole[msg.Role] = append(byRole[msg.Role], msg)
	}
	for _, role := range roleOrder {
		fmt.Fprintf(&b, "  %s:\n", role)
		messages := byRole[role]
		for i, msg := range messages {
			fmt.Fprintf(&b, "    - message: %d\n", i+1)
			if msg.CacheControl != nil {
				fmt.Fprintf(&b, "      cache_control: %s", msg.CacheControl.Type)
				if msg.CacheControl.TTL != "" {
					fmt.Fprintf(&b, " ttl=%s", msg.CacheControl.TTL)
				}
				fmt.Fprintln(&b)
			}
			fmt.Fprintln(&b, "      content: |")
			content := msg.Content
			if content == "" {
				fmt.Fprintln(&b, "        ")
				continue
			}
			for _, line := range strings.Split(content, "\n") {
				fmt.Fprintf(&b, "        %s\n", line)
			}
		}
	}

	return b.Bytes(), nil
}

func formatExtraParamTextValue(value any) (string, error) {
	if s, ok := value.(string); ok {
		return s, nil
	}

	b, err := json.Marshal(value)
	if err != nil {
		return "", err
	}
	return string(b), nil
}