package llm

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"os"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"

	"gitea.maximumdirect.net/eric/notarius/internal/core/artifacts"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/promptkit"
)

const promptKitProviderName = "promptkit"

type PromptKitLocalBackendConfig struct {
	Endpoint         string
	ConcurrencyLimit int
}

type PromptKitClientConfig struct {
	ProfileDir      string
	ProfileFile     string
	LocalBackend    *PromptKitLocalBackendConfig
	Assets          *AssetRegistry
	Timeout         time.Duration
	HTTPClient      *http.Client
	EngineOptions   []promptkit.Option
	Recorder        *LLMProfileRecorder
	ReasoningEffort *string
}

type PromptKitClient struct {
	engine                *promptkit.Engine
	recorder              *LLMProfileRecorder
	profileDir            string
	profileFile           string
	localEndpoint         string
	fallbackProfileDigest string
	reasoningEffort       *string
}

type LLMProfileRecorder struct {
	mu       sync.Mutex
	profiles map[string]artifacts.LLMProfileManifest
}

var _ contracts.StructuredLLMClient = (*PromptKitClient)(nil)
var _ contracts.LLMProfileManifestProvider = (*PromptKitClient)(nil)

func PromptKitLocalBackendOption(cfg PromptKitLocalBackendConfig) promptkit.Option {
	return promptkit.WithBackend(promptkit.LocalBackend(cfg.Endpoint, cfg.ConcurrencyLimit))
}

func NewPromptKitClient(cfg PromptKitClientConfig) (*PromptKitClient, error) {
	if cfg.Assets == nil {
		return nil, fmt.Errorf("PromptKit client assets must not be nil")
	}
	profileSource, profileOptions, err := promptKitProfileSourceEngineOptions(PromptKitProfileSourceConfig{
		ProfileDir:   cfg.ProfileDir,
		ProfileFile:  cfg.ProfileFile,
		LocalBackend: cfg.LocalBackend,
	})
	if err != nil {
		return nil, err
	}
	options, fallbackProfileDigest, err := cfg.Assets.promptKitOptions()
	if err != nil {
		return nil, err
	}
	options = append(options, profileOptions...)
	options = append(options, cfg.EngineOptions...)

	engine, err := promptkit.NewEngine(promptkit.Config{
		ProfileDir: profileSource.ProfileDir,
		Timeout:    cfg.Timeout,
		HTTPClient: cfg.HTTPClient,
	}, options...)
	if err != nil {
		return nil, fmt.Errorf("create PromptKit engine: %w", err)
	}
	recorder := cfg.Recorder
	if recorder == nil {
		recorder = NewLLMProfileRecorder()
	}
	var reasoningEffort *string
	if cfg.ReasoningEffort != nil {
		value := *cfg.ReasoningEffort
		reasoningEffort = &value
	}
	return &PromptKitClient{
		engine:                engine,
		recorder:              recorder,
		profileDir:            profileSource.ProfileDir,
		profileFile:           profileSource.ProfileFile,
		localEndpoint:         profileSource.localEndpoint(),
		fallbackProfileDigest: fallbackProfileDigest,
		reasoningEffort:       reasoningEffort,
	}, nil
}

func (c *PromptKitClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
	if c == nil {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("PromptKit client must not be nil")
	}
	if c.engine == nil {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("PromptKit client engine must not be nil")
	}
	if err := validateOutputTarget(out); err != nil {
		return contracts.StructuredCompletionResponse{}, err
	}
	if req.StructuredOutputRepairAttempts != nil && (*req.StructuredOutputRepairAttempts < 0 || *req.StructuredOutputRepairAttempts > 3) {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("structured output repair attempts must be between zero and three")
	}
	appendedMessages, err := promptKitCorrectionMessages(req.Correction)
	if err != nil {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("structured completion correction: %w", err)
	}
	promptID := strings.TrimSpace(req.PromptID)
	if promptID == "" {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("structured completion prompt_id must not be empty")
	}
	sessionID := strings.TrimSpace(req.SessionID)
	var execution *promptkit.ExecutionTargetOverride
	if c.reasoningEffort != nil {
		reasoningEffort := *c.reasoningEffort
		execution = &promptkit.ExecutionTargetOverride{
			ReasoningEffort: &reasoningEffort,
		}
	}

	runReq := promptkit.RunRequest{
		PromptID:         promptID,
		PromptVersion:    strings.TrimSpace(req.PromptVersion),
		ProfileID:        strings.TrimSpace(req.ProfileID),
		SessionID:        sessionID,
		Inputs:           promptKitInputs(req.Inputs),
		Vars:             promptKitVars(req, sessionID),
		Execution:        execution,
		AppendedMessages: appendedMessages,
	}
	if req.StructuredOutputRepairAttempts != nil {
		inspection, err := c.engine.InspectPrompt(ctx, promptID, strings.TrimSpace(req.PromptVersion))
		if err != nil {
			if ctxErr := ctx.Err(); ctxErr != nil {
				return contracts.StructuredCompletionResponse{}, ctxErr
			}
			return contracts.StructuredCompletionResponse{}, fmt.Errorf("inspect PromptKit prompt %q: %v", promptID, redactPromptKitError(err))
		}
		contract := inspection.OutputContract
		contract.RepairAttempts = *req.StructuredOutputRepairAttempts
		runReq.Validation = &contract
	}
	prepared, err := c.engine.PrepareExecution(ctx, runReq)
	if err != nil {
		if ctxErr := ctx.Err(); ctxErr != nil {
			return contracts.StructuredCompletionResponse{}, ctxErr
		}
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("prepare PromptKit prompt %q: %v", promptID, redactPromptKitError(err))
	}
	defer prepared.Discard()
	preparedDetails := prepared.Details()
	result, err := c.engine.RunPrepared(ctx, prepared)
	if err != nil {
		if ctxErr := ctx.Err(); ctxErr != nil {
			return contracts.StructuredCompletionResponse{}, ctxErr
		}
		if errors.Is(err, promptkit.ErrCapacityExceeded) {
			var capacityErr *promptkit.CapacityError
			if errors.As(err, &capacityErr) && strings.TrimSpace(capacityErr.BackendID) != "" {
				return contracts.StructuredCompletionResponse{}, fmt.Errorf(
					"run PromptKit prompt %q on backend %q: %w: %v",
					promptID,
					strings.TrimSpace(capacityErr.BackendID),
					contracts.ErrLLMCapacityExceeded,
					redactPromptKitError(err),
				)
			}
			return contracts.StructuredCompletionResponse{}, fmt.Errorf(
				"run PromptKit prompt %q: %w: %v",
				promptID,
				contracts.ErrLLMCapacityExceeded,
				redactPromptKitError(err),
			)
		}
		if errors.Is(err, promptkit.ErrLLMGenerate) {
			var generationErr *promptkit.GenerationError
			status := 0
			response := contracts.StructuredCompletionResponse{Debug: &contracts.LLMDebugMaterial{Prompt: promptKitDebugPrompt(&preparedDetails)}}
			if errors.As(err, &generationErr) {
				status = generationErr.StatusCode()
				response.Debug.Response = promptKitDebugGenerationError(&preparedDetails, generationErr)
			}
			return response, contracts.NewLLMGenerationError(status, fmt.Sprintf("run PromptKit prompt %q: %v", promptID, redactPromptKitError(err)))
		}
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("run PromptKit prompt %q: %v", promptID, redactPromptKitError(err))
	}
	if result == nil {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("run PromptKit prompt %q: %w: empty result", promptID, contracts.ErrInvalidStructuredOutput)
	}
	response := c.responseFromResult(result, &preparedDetails)
	if result.Validation.Status == promptkit.ValidationFailed || !result.Validation.IsValid {
		return response, fmt.Errorf("run PromptKit prompt %q: %w: validation failed: %s", promptID, contracts.ErrInvalidStructuredOutput, strings.Join(result.Validation.Errors, "; "))
	}
	if len(strings.TrimSpace(string(response.Content))) == 0 {
		return response, fmt.Errorf("run PromptKit prompt %q: %w: empty structured output", promptID, contracts.ErrInvalidStructuredOutput)
	}
	if err := json.Unmarshal(response.Content, out); err != nil {
		return response, fmt.Errorf("decode PromptKit structured output for prompt %q: %w: %w", promptID, contracts.ErrInvalidStructuredOutput, err)
	}
	return response, nil
}

func promptKitCorrectionMessages(correction *contracts.SemanticCorrection) ([]promptkit.RenderedMessage, error) {
	owned, err := contracts.CloneSemanticCorrection(correction)
	if err != nil {
		return nil, err
	}
	if owned == nil {
		return nil, nil
	}
	return []promptkit.RenderedMessage{
		{Role: promptkit.RoleAssistant, Content: string(owned.AssistantResponse)},
		{Role: promptkit.RoleUser, Content: owned.UserGuidance},
	}, nil
}

func promptKitDebugGenerationError(prepared *promptkit.PreparedRun, generationErr *promptkit.GenerationError) *contracts.LLMDebugResponse {
	if generationErr == nil {
		return nil
	}
	var secrets []string
	if prepared != nil && strings.TrimSpace(prepared.EffectiveModelParams.APIKeyEnv) != "" {
		if value, ok := os.LookupEnv(prepared.EffectiveModelParams.APIKeyEnv); ok {
			secrets = append(secrets, value)
		}
	}
	redact := func(value string) string {
		return bearerTokenPattern.ReplaceAllString(RedactSecrets(value, secrets), "Bearer "+secretReplacement)
	}
	return &contracts.LLMDebugResponse{ProviderError: &contracts.LLMDebugProviderError{
		StatusCode: generationErr.StatusCode(),
		Code:       redact(generationErr.ProviderCode()),
		Type:       redact(generationErr.ProviderType()),
		Message:    redact(generationErr.ProviderMessage()),
	}}
}

func (c *PromptKitClient) responseFromResult(result *promptkit.RunResult, prepared *promptkit.PreparedRun) contracts.StructuredCompletionResponse {
	content := result.Artifact.Body
	if len(content) == 0 {
		content = []byte(result.RawOutput)
	}
	profile := artifacts.LLMProfileManifest{
		ID:              strings.TrimSpace(result.SelectedProfileID),
		Provider:        promptKitProviderName,
		Model:           firstNonEmpty(result.ModelName, result.EffectiveModelParams.Model),
		BackendID:       strings.TrimSpace(result.SelectedBackendID),
		ReasoningEffort: strings.TrimSpace(result.EffectiveModelParams.ReasoningEffort),
	}
	if c.recorder != nil {
		c.recorder.Record(profile)
	}
	return contracts.StructuredCompletionResponse{
		Content:          append(json.RawMessage(nil), content...),
		Provider:         profile.Provider,
		Model:            profile.Model,
		ProfileID:        profile.ID,
		PromptTokens:     result.Usage.PromptTokens,
		CompletionTokens: result.Usage.CompletionTokens,
		TotalTokens:      result.Usage.TotalTokens,
		RepairAttempts:   result.Validation.RepairAttempts,
		Debug:            promptKitDebugMaterial(prepared, result),
	}
}

func promptKitDebugMaterial(prepared *promptkit.PreparedRun, result *promptkit.RunResult) *contracts.LLMDebugMaterial {
	material := &contracts.LLMDebugMaterial{}
	if prepared != nil {
		material.Prompt = promptKitDebugPrompt(prepared)
	}
	if result != nil {
		material.Response = promptKitDebugResponse(result)
	}
	if material.Prompt == nil && material.Response == nil {
		return nil
	}
	return material
}

func promptKitDebugPrompt(prepared *promptkit.PreparedRun) *contracts.LLMDebugPrompt {
	if prepared == nil {
		return nil
	}
	messages := make([]contracts.LLMDebugMessage, 0, len(prepared.Messages))
	for _, message := range prepared.Messages {
		messages = append(messages, contracts.LLMDebugMessage{
			Role:         message.Role,
			Content:      message.Content,
			CacheControl: jsonObject(message.CacheControl),
		})
	}
	return &contracts.LLMDebugPrompt{
		PromptID:             prepared.PromptID,
		PromptVersion:        prepared.PromptVersion,
		PromptHash:           prepared.PromptHash,
		SelectedProfileID:    prepared.SelectedProfileID,
		SelectedBackendID:    prepared.SelectedBackendID,
		SessionID:            prepared.SessionID,
		RenderedPromptHash:   prepared.RenderedPromptHash,
		Messages:             messages,
		EffectiveModelParams: jsonObject(prepared.EffectiveModelParams),
		OutputContract:       jsonObject(prepared.OutputContract),
		StructuredOutput:     jsonObject(prepared.StructuredOutput),
		InputHashes:          cloneStringMap(prepared.InputHashes),
	}
}

func promptKitDebugResponse(result *promptkit.RunResult) *contracts.LLMDebugResponse {
	if result == nil {
		return nil
	}
	content := result.RawOutput
	if content == "" {
		content = string(result.Artifact.Body)
	}
	return &contracts.LLMDebugResponse{
		Content:              content,
		RunID:                result.RunID,
		PromptID:             result.PromptID,
		PromptVersion:        result.PromptVersion,
		PromptHash:           result.PromptHash,
		RenderedPromptHash:   result.RenderedPromptHash,
		SelectedProfileID:    result.SelectedProfileID,
		ModelName:            result.ModelName,
		Endpoint:             result.Endpoint,
		EffectiveModelParams: jsonObject(result.EffectiveModelParams),
		InputHashes:          cloneStringMap(result.InputHashes),
		Validation:           jsonObject(result.Validation),
		Usage: contracts.LLMDebugUsage{
			PromptTokens:     result.Usage.PromptTokens,
			CompletionTokens: result.Usage.CompletionTokens,
			TotalTokens:      result.Usage.TotalTokens,
			CachedTokens:     result.Usage.CachedTokens,
			CacheWriteTokens: result.Usage.CacheWriteTokens,
		},
	}
}

func jsonObject(value any) map[string]any {
	if value == nil {
		return nil
	}
	data, err := json.Marshal(value)
	if err != nil || string(data) == "null" {
		return nil
	}
	var out map[string]any
	if err := json.Unmarshal(data, &out); err != nil {
		return nil
	}
	if len(out) == 0 {
		return nil
	}
	return out
}

func cloneStringMap(values map[string]string) map[string]string {
	if len(values) == 0 {
		return nil
	}
	out := make(map[string]string, len(values))
	for key, value := range values {
		out[key] = value
	}
	return out
}

func (c *PromptKitClient) LLMProfileManifests() []artifacts.LLMProfileManifest {
	if c == nil || c.recorder == nil {
		return nil
	}
	return c.recorder.Manifests()
}

func (c *PromptKitClient) LLMCheckpointFingerprints() ([]CheckpointFingerprint, error) {
	if c == nil {
		return nil, nil
	}
	fingerprint, err := promptKitProfileFingerprint(c.profileDir, c.profileFile, c.fallbackProfileDigest)
	if err != nil {
		return nil, err
	}
	fingerprints := []CheckpointFingerprint{fingerprint}
	if c.localEndpoint != "" {
		fingerprints = append(fingerprints, promptKitLocalBackendFingerprint(c.localEndpoint))
	}
	return fingerprints, nil
}

func NewLLMProfileRecorder() *LLMProfileRecorder {
	return &LLMProfileRecorder{profiles: map[string]artifacts.LLMProfileManifest{}}
}

func (r *LLMProfileRecorder) Record(profile artifacts.LLMProfileManifest) {
	if r == nil {
		return
	}
	profile = profile.Normalized()
	key := profile.IdentityKey()
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.profiles == nil {
		r.profiles = map[string]artifacts.LLMProfileManifest{}
	}
	r.profiles[key] = profile
}

func (r *LLMProfileRecorder) Manifests() []artifacts.LLMProfileManifest {
	if r == nil {
		return nil
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	if len(r.profiles) == 0 {
		return nil
	}
	keys := make([]string, 0, len(r.profiles))
	for key := range r.profiles {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	out := make([]artifacts.LLMProfileManifest, 0, len(keys))
	for _, key := range keys {
		out = append(out, r.profiles[key])
	}
	return out
}

func promptKitInputs(inputs contracts.LLMInputSet) map[string]promptkit.ArtifactRef {
	if len(inputs) == 0 {
		return nil
	}
	out := make(map[string]promptkit.ArtifactRef, len(inputs))
	for key, material := range inputs {
		name := strings.TrimSpace(key)
		if name == "" {
			name = strings.TrimSpace(material.Name)
		}
		if name == "" {
			continue
		}
		body := string(material.Content)
		if body == "" {
			body = " "
		}
		if origin := strings.TrimSpace(material.OriginURI); origin != "" {
			out[name] = promptkit.InlineWithURI(origin, body)
		} else {
			out[name] = promptkit.Inline(body)
		}
	}
	return out
}

func promptKitVars(req contracts.StructuredCompletionRequest, sessionID string) map[string]string {
	vars := make(map[string]string, len(req.Vars)+1)
	for key, value := range req.Vars {
		name := strings.TrimSpace(key)
		if name == "" || value == nil {
			continue
		}
		vars[name] = fmt.Sprint(value)
	}
	if sessionID != "" {
		vars["session_id"] = sessionID
	}
	if len(vars) == 0 {
		return nil
	}
	return vars
}

var bearerTokenPattern = regexp.MustCompile(`(?i)Bearer\s+[A-Za-z0-9._~+/=-]+`)

func redactPromptKitError(err error) error {
	if err == nil {
		return nil
	}
	return sanitizedProviderDiagnostic{message: bearerTokenPattern.ReplaceAllString(err.Error(), "Bearer "+secretReplacement)}
}

type sanitizedProviderDiagnostic struct {
	message string
}

func (e sanitizedProviderDiagnostic) Error() string {
	return e.message
}
