package appendorder

import (
	"context"
	"encoding/json"
	"fmt"

	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
)

const Key = "appendorder"

var _ contracts.Merger = (*Merger)(nil)

type Merger struct{}

func New() *Merger {
	return &Merger{}
}

func (m *Merger) Key() string {
	return Key
}

func (m *Merger) Merge(ctx context.Context, req contracts.MergeRequest) (contracts.MergeResult, error) {
	if m == nil {
		return contracts.MergeResult{}, mergerErrorf("merger must not be nil")
	}
	if ctx == nil {
		return contracts.MergeResult{}, mergerErrorf("context must not be nil")
	}
	if err := ctx.Err(); err != nil {
		return contracts.MergeResult{}, mergerErrorf("context error before merge: %w", err)
	}

	if len(req.ExtractOutputs) == 1 {
		payload := cloneRawPayload(req.ExtractOutputs[0].Payload)
		return contracts.MergeResult{
			Output: contracts.MergeOutput{
				LaneID:    req.LaneID,
				MergerKey: Key,
				SourceID:  req.ExtractOutputs[0].SourceID,
				Schema:    req.ExtractOutputs[0].Schema,
				Payload:   payload,
			},
		}, nil
	}

	content, err := orderedContent(req.ExtractOutputs)
	if err != nil {
		return contracts.MergeResult{}, err
	}
	return contracts.MergeResult{
		Output: contracts.MergeOutput{
			LaneID:    req.LaneID,
			MergerKey: Key,
			SourceID:  sourceID(req.ExtractOutputs),
			Payload: contracts.RawPayload{
				Content:   content,
				MediaType: "application/json",
			},
		},
	}, nil
}

func ModuleSpec() pipeline.ModuleSpec {
	return pipeline.ModuleSpec{
		Key:      Key,
		Stage:    pipeline.StageMerge,
		Provides: []string{"merged"},
	}
}

func Register(registry *pipeline.MergerRegistry) error {
	return registry.RegisterWithSpec(ModuleSpec(), func() (contracts.Merger, error) {
		return New(), nil
	})
}

func orderedContent(outputs []contracts.ExtractOutput) ([]byte, error) {
	items := make([]map[string]any, 0, len(outputs))
	for _, output := range outputs {
		item := map[string]any{
			"chunk_id":    output.ChunkID,
			"chunk_index": output.ChunkIndex,
			"media_type":  output.Payload.MediaType,
		}
		if json.Valid(output.Payload.Content) {
			item["content"] = json.RawMessage(append([]byte(nil), output.Payload.Content...))
		} else {
			item["content"] = string(output.Payload.Content)
		}
		items = append(items, item)
	}
	content, err := json.Marshal(struct {
		Outputs []map[string]any `json:"outputs"`
	}{Outputs: items})
	if err != nil {
		return nil, mergerErrorf("encode merged output: %w", err)
	}
	return content, nil
}

func sourceID(outputs []contracts.ExtractOutput) string {
	for _, output := range outputs {
		if output.SourceID != "" {
			return output.SourceID
		}
	}
	return ""
}

func cloneRawPayload(payload contracts.RawPayload) contracts.RawPayload {
	return contracts.RawPayload{
		Content:   append([]byte(nil), payload.Content...),
		MediaType: payload.MediaType,
		Metadata:  cloneMetadata(payload.Metadata),
		Warnings:  append([]contracts.Warning(nil), payload.Warnings...),
	}
}

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

func mergerErrorf(format string, args ...any) error {
	return fmt.Errorf("appendorder merger: "+format, args...)
}