package llm

import (
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"gitea.maximumdirect.net/eric/audita/internal/core/redaction"
)

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 {
	return &DiagnosticsWriter{dir: dir, secrets: append([]string(nil), secrets...)}
}

// 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 {
	redacted, err := redaction.New(w.secrets).JSON(payload)
	if err != nil {
		return fmt.Errorf("marshal json artifact %q: %w", filepath.Base(path), err)
	}
	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
}