package llm

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"
)

const redactedSecret = "[REDACTED]"

var stageSanitizer = regexp.MustCompile(`[^a-zA-Z0-9_.-]+`)

// InteractionArtifacts contains file paths for written prompt/response
// diagnostics.
type InteractionArtifacts struct {
	RequestMetadataPath string
	RequestPayloadPath  string
	ResponsePayloadPath string
	ErrorPayloadPath    string
}

// DiagnosticsWriter writes generic machine-readable LLM interaction artifacts.
type DiagnosticsWriter struct {
	dir     string
	secrets []string
}

// NewDiagnosticsWriter creates a diagnostics writer rooted at dir.
func NewDiagnosticsWriter(dir string, secrets []string) *DiagnosticsWriter {
	filtered := make([]string, 0, len(secrets))
	for _, secret := range secrets {
		secret = strings.TrimSpace(secret)
		if secret != "" {
			filtered = append(filtered, secret)
		}
	}

	return &DiagnosticsWriter{
		dir:     dir,
		secrets: filtered,
	}
}

// WriteInteraction writes request/response/error artifacts for a generic LLM
// interaction stage.
func (w *DiagnosticsWriter) WriteInteraction(stage string, requestMetadata any, requestPayload any, responsePayload any, errorPayload any) (InteractionArtifacts, error) {
	if strings.TrimSpace(w.dir) == "" {
		return InteractionArtifacts{}, fmt.Errorf("diagnostics directory must not be empty")
	}
	if err := os.MkdirAll(w.dir, 0o755); err != nil {
		return InteractionArtifacts{}, fmt.Errorf("create diagnostics directory: %w", err)
	}

	base := sanitizeStageName(stage)
	paths := InteractionArtifacts{
		RequestMetadataPath: filepath.Join(w.dir, base+"-request-metadata.json"),
		RequestPayloadPath:  filepath.Join(w.dir, base+"-request-payload.json"),
		ResponsePayloadPath: filepath.Join(w.dir, base+"-response-payload.json"),
	}

	if err := w.writeJSON(paths.RequestMetadataPath, requestMetadata); err != nil {
		return InteractionArtifacts{}, err
	}
	if err := w.writeJSON(paths.RequestPayloadPath, requestPayload); err != nil {
		return InteractionArtifacts{}, err
	}
	if err := w.writeJSON(paths.ResponsePayloadPath, responsePayload); err != nil {
		return InteractionArtifacts{}, err
	}
	if errorPayload != nil {
		paths.ErrorPayloadPath = filepath.Join(w.dir, base+"-error-payload.json")
		if err := w.writeJSON(paths.ErrorPayloadPath, errorPayload); err != nil {
			return InteractionArtifacts{}, err
		}
	}

	return paths, nil
}

func sanitizeStageName(stage string) string {
	stage = strings.TrimSpace(stage)
	if stage == "" {
		return "llm"
	}
	safe := stageSanitizer.ReplaceAllString(stage, "_")
	safe = strings.Trim(safe, "._-")
	if safe == "" {
		return "llm"
	}
	return safe
}

func (w *DiagnosticsWriter) writeJSON(path string, payload any) error {
	raw, err := json.MarshalIndent(payload, "", "  ")
	if err != nil {
		return fmt.Errorf("marshal json artifact %q: %w", filepath.Base(path), err)
	}
	redacted := redactJSONSecrets(raw, w.secrets)
	if !json.Valid(redacted) {
		return fmt.Errorf("redacted json artifact %q is invalid", filepath.Base(path))
	}
	redacted = append(redacted, '\n')
	if err := os.WriteFile(path, redacted, 0o644); err != nil {
		return fmt.Errorf("write json artifact %q: %w", filepath.Base(path), err)
	}
	return nil
}

func redactJSONSecrets(raw []byte, secrets []string) []byte {
	if len(secrets) == 0 {
		return raw
	}

	result := string(raw)
	for _, secret := range secrets {
		if secret == "" {
			continue
		}
		result = strings.ReplaceAll(result, "Bearer "+secret, "Bearer "+redactedSecret)
		result = strings.ReplaceAll(result, secret, redactedSecret)
	}
	return []byte(result)
}