package llm

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"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/scriptorium"
)

const scriptoriumProviderName = "scriptorium"

type ScriptoriumClientConfig struct {
	ProfileDir    string
	ProfileFile   string
	Assets        *AssetRegistry
	Timeout       time.Duration
	HTTPClient    *http.Client
	EngineOptions []scriptorium.Option
	Recorder      *LLMProfileRecorder
}

type ScriptoriumClient struct {
	engine   *scriptorium.Engine
	recorder *LLMProfileRecorder
}

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

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

func NewScriptoriumClient(cfg ScriptoriumClientConfig) (*ScriptoriumClient, error) {
	if cfg.Assets == nil {
		return nil, fmt.Errorf("scriptorium client assets must not be nil")
	}
	if strings.TrimSpace(cfg.ProfileDir) != "" && strings.TrimSpace(cfg.ProfileFile) != "" {
		return nil, fmt.Errorf("scriptorium profile_dir and profile_file are mutually exclusive")
	}
	options, err := cfg.Assets.ScriptoriumOptions()
	if err != nil {
		return nil, err
	}
	if profileFile := strings.TrimSpace(cfg.ProfileFile); profileFile != "" {
		options = append(options, scriptorium.WithProfileFile(profileFile))
	}
	options = append(options, cfg.EngineOptions...)

	engine, err := scriptorium.NewEngine(scriptorium.Config{
		ProfileDir: strings.TrimSpace(cfg.ProfileDir),
		Timeout:    cfg.Timeout,
		HTTPClient: cfg.HTTPClient,
	}, options...)
	if err != nil {
		return nil, fmt.Errorf("create Scriptorium engine: %w", err)
	}
	recorder := cfg.Recorder
	if recorder == nil {
		recorder = NewLLMProfileRecorder()
	}
	return &ScriptoriumClient{
		engine:   engine,
		recorder: recorder,
	}, nil
}

func (c *ScriptoriumClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
	if c == nil {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("scriptorium client must not be nil")
	}
	if c.engine == nil {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("scriptorium client engine must not be nil")
	}
	if err := validateOutputTarget(out); err != nil {
		return contracts.StructuredCompletionResponse{}, err
	}
	promptID := strings.TrimSpace(req.PromptID)
	if promptID == "" {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("structured completion prompt_id must not be empty")
	}

	runReq := scriptorium.RunRequest{
		PromptID:      promptID,
		PromptVersion: strings.TrimSpace(req.PromptVersion),
		ProfileID:     strings.TrimSpace(req.ProfileID),
		Inputs:        scriptoriumInputs(req.Inputs),
		Vars:          scriptoriumVars(req),
		Metadata:      scriptoriumMetadata(req),
	}
	prepared, err := c.engine.Prepare(ctx, runReq)
	if err != nil {
		if ctxErr := ctx.Err(); ctxErr != nil {
			return contracts.StructuredCompletionResponse{}, ctxErr
		}
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("prepare Scriptorium prompt %q: %w", promptID, redactScriptoriumError(err))
	}
	result, err := c.engine.Run(ctx, runReq)
	if err != nil {
		if ctxErr := ctx.Err(); ctxErr != nil {
			return contracts.StructuredCompletionResponse{}, ctxErr
		}
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("run Scriptorium prompt %q: %w", promptID, redactScriptoriumError(err))
	}
	if result == nil {
		return contracts.StructuredCompletionResponse{}, fmt.Errorf("run Scriptorium prompt %q: empty result", promptID)
	}
	response := c.responseFromResult(result, prepared)
	if result.Validation.Status == scriptorium.ValidationFailed || !result.Validation.IsValid {
		return response, fmt.Errorf("run Scriptorium prompt %q: validation failed: %s", promptID, strings.Join(result.Validation.Errors, "; "))
	}
	if len(strings.TrimSpace(string(response.Content))) == 0 {
		return response, fmt.Errorf("run Scriptorium prompt %q: empty structured output", promptID)
	}
	if err := json.Unmarshal(response.Content, out); err != nil {
		return response, fmt.Errorf("decode Scriptorium structured output for prompt %q: %w", promptID, err)
	}
	return response, nil
}

func (c *ScriptoriumClient) responseFromResult(result *scriptorium.RunResult, prepared *scriptorium.PreparedRun) contracts.StructuredCompletionResponse {
	content := result.Artifact.Body
	if len(content) == 0 {
		content = []byte(result.RawOutput)
	}
	profile := artifacts.LLMProfileManifest{
		ID:       strings.TrimSpace(result.SelectedProfileID),
		Provider: scriptoriumProviderName,
		Model:    firstNonEmpty(result.ModelName, result.EffectiveModelParams.Model),
	}
	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,
		Debug:            scriptoriumDebugMaterial(prepared, result),
	}
}

func scriptoriumDebugMaterial(prepared *scriptorium.PreparedRun, result *scriptorium.RunResult) *contracts.LLMDebugMaterial {
	material := &contracts.LLMDebugMaterial{}
	if prepared != nil {
		material.Prompt = scriptoriumDebugPrompt(prepared)
	}
	if result != nil {
		material.Response = scriptoriumDebugResponse(result)
	}
	if material.Prompt == nil && material.Response == nil {
		return nil
	}
	return material
}

func scriptoriumDebugPrompt(prepared *scriptorium.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,
		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 scriptoriumDebugResponse(result *scriptorium.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 *ScriptoriumClient) LLMProfileManifests() []artifacts.LLMProfileManifest {
	if c == nil || c.recorder == nil {
		return nil
	}
	return c.recorder.Manifests()
}

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

func (r *LLMProfileRecorder) Record(profile artifacts.LLMProfileManifest) {
	if r == nil {
		return
	}
	profile.ID = strings.TrimSpace(profile.ID)
	profile.Provider = strings.TrimSpace(profile.Provider)
	profile.Model = strings.TrimSpace(profile.Model)
	key := profile.ID + "\x00" + profile.Provider + "\x00" + profile.Model
	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 scriptoriumInputs(inputs contracts.LLMInputSet) map[string]scriptorium.ArtifactRef {
	if len(inputs) == 0 {
		return nil
	}
	out := make(map[string]scriptorium.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] = scriptorium.InlineWithURI(origin, body)
		} else {
			out[name] = scriptorium.Inline(body)
		}
	}
	return out
}

func scriptoriumVars(req contracts.StructuredCompletionRequest) 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 := strings.TrimSpace(req.SessionID); sessionID != "" {
		vars["session_id"] = sessionID
	}
	if len(vars) == 0 {
		return nil
	}
	return vars
}

func scriptoriumMetadata(req contracts.StructuredCompletionRequest) map[string]string {
	metadata := map[string]string{}
	if stageName := strings.TrimSpace(req.StageName); stageName != "" {
		metadata["stage_name"] = stageName
	}
	if len(metadata) == 0 {
		return nil
	}
	return metadata
}

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

func redactScriptoriumError(err error) error {
	if err == nil {
		return nil
	}
	return redactedProviderError{err: err}
}

type redactedProviderError struct {
	err error
}

func (e redactedProviderError) Error() string {
	return bearerTokenPattern.ReplaceAllString(e.err.Error(), "Bearer "+secretReplacement)
}

func (e redactedProviderError) Unwrap() error {
	return e.err
}