package spells

import (
	"bytes"
	"context"
	"fmt"

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

const Key = "dnd/spells"
const ArtifactType = "dnd.spell_cast"
const SchemaVersion = "v1"

var requiredCapabilities = []string{
	"chunks",
	"source.transcript",
}

var providedCapabilities = []string{
	"dnd.spell_casts",
}

var referenceSlotDescriptions = shared.ReferenceSlotDescriptions{
	Glossary: "Optional campaign glossary reference material used only for disambiguation.",
	Party:    "Optional party roster reference material used only for disambiguation.",
	Players:  "Optional player list reference material used only for disambiguation.",
	Roster:   "Deprecated alias for party roster reference material used only for disambiguation.",
}

var _ contracts.Extractor[dnd.SpellList] = (*Extractor)(nil)

type Options struct{}

type Extractor struct {
	llm contracts.StructuredLLMClient
}

type rawAdapter struct {
	extractor *Extractor
	codec     RawAdapterCodec
}

type RawAdapterCodec interface {
	contracts.ArtifactCodec[dnd.SpellList]
	EncodeCandidate(dnd.SpellList) ([]byte, error)
}

var _ contracts.LegacyRawExtractor = (*rawAdapter)(nil)

func New(llmClient contracts.StructuredLLMClient, _ Options) (*Extractor, error) {
	if llmClient == nil {
		return nil, extractorErrorf("LLM client must not be nil")
	}
	return &Extractor{llm: llmClient}, nil
}

func (e *Extractor) Key() string {
	return Key
}

func (e *Extractor) ReferenceSlots() []contracts.ReferenceSlot {
	return shared.ReferenceSlots(referenceSlotDescriptions)
}

func (e *Extractor) ManifestMetadata() map[string]any {
	promptSHA, err := scriptoriumPromptMetadata()
	if err != nil {
		promptSHA = ""
	}
	metadata := map[string]any{
		"prompt_id":            PromptID,
		"prompt_version":       SchemaVersion,
		"prompt_sha256":        promptSHA,
		"response_schema_key":  string(ResponseSchemaKey),
		"response_schema_id":   ResponseSchemaID,
		"response_schema_name": ResponseSchemaName,
	}
	if schema, err := loadResponseSchema(); err == nil {
		metadata["response_schema_version"] = schema.Version
		metadata["response_schema_sha256"] = schema.SHA256
	}
	return metadata
}

func (e *Extractor) Extract(ctx context.Context, req contracts.TypedExtractionRequest) (contracts.TypedExtractionResult[dnd.SpellList], error) {
	if e == nil {
		return contracts.TypedExtractionResult[dnd.SpellList]{}, extractorErrorf("extractor must not be nil")
	}
	if e.llm == nil {
		return contracts.TypedExtractionResult[dnd.SpellList]{}, extractorErrorf("LLM client must not be nil")
	}
	if ctx == nil {
		return contracts.TypedExtractionResult[dnd.SpellList]{}, extractorErrorf("context must not be nil")
	}
	if err := ctx.Err(); err != nil {
		return contracts.TypedExtractionResult[dnd.SpellList]{}, extractorErrorf("context error before extraction: %w", err)
	}
	if req.Source == nil {
		return contracts.TypedExtractionResult[dnd.SpellList]{}, extractorErrorf("source must not be nil")
	}
	if req.Chunk == nil {
		return contracts.TypedExtractionResult[dnd.SpellList]{}, extractorErrorf("chunk must not be nil")
	}
	if len(req.Chunk.Units) == 0 {
		return contracts.TypedExtractionResult[dnd.SpellList]{}, extractorErrorf("chunk %q units must not be empty", req.Chunk.ID)
	}
	sourceInput, err := chunkSourceInput(req)
	if err != nil {
		return contracts.TypedExtractionResult[dnd.SpellList]{}, err
	}

	var response extractionResponse
	if _, err := e.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
		StageName:     Key,
		PromptID:      PromptID,
		PromptVersion: SchemaVersion,
		ProfileID:     req.LLMProfile,
		SessionID:     req.SessionID,
		Inputs:        shared.PromptInputs(sourceInput, req.References),
	}, &response); err != nil {
		return contracts.TypedExtractionResult[dnd.SpellList]{}, extractorErrorf("complete structured output: %w", err)
	}
	canonicalizeResponse(&response)
	return contracts.TypedExtractionResult[dnd.SpellList]{Value: canonicalSpellList(response, req.Source.ID)}, nil
}

func chunkSourceInput(req contracts.TypedExtractionRequest) (contracts.LLMInputMaterial, error) {
	material := req.SourceInput.Clone()
	if len(material.Content) == 0 {
		material = contracts.NewLLMInputMaterial("source", req.Chunk.MediaType, req.Chunk.Content, "", "")
	}
	if !bytes.Equal(material.Content, req.Chunk.Content) {
		return contracts.LLMInputMaterial{}, extractorErrorf("source input must match chunk %q content", req.Chunk.ID)
	}
	if material.Name == "" {
		material.Name = "source"
	}
	if material.MediaType == "" {
		material.MediaType = req.Chunk.MediaType
	}
	if material.SizeBytes == 0 {
		material.SizeBytes = int64(len(material.Content))
	}
	return material, nil
}

func ModuleSpec() pipeline.ModuleSpec {
	return pipeline.ModuleSpec{
		Key:            Key,
		Stage:          pipeline.StageExtract,
		Requires:       append([]string(nil), requiredCapabilities...),
		Provides:       append([]string(nil), providedCapabilities...),
		ArtifactKind:   dnd.SpellListKind,
		ReferenceSlots: shared.ReferenceSlots(referenceSlotDescriptions),
	}
}

func Register(registry *pipeline.ExtractorRegistry) error {
	return pipeline.RegisterExtractorBuilder(registry, ModuleSpec(), validateOptions, func(request pipeline.BuildRequest) (contracts.Extractor[dnd.SpellList], error) {
		options, err := DecodeOptions(request.Options)
		if err != nil {
			return nil, err
		}
		return New(request.Dependencies.LLM, options)
	})
}

// RegisterWithRawAdapter keeps existing raw downstream implementations usable
// while the extractor itself produces the canonical typed artifact.
func RegisterWithRawAdapter(registry *pipeline.ExtractorRegistry, codec RawAdapterCodec) error {
	if codec == nil {
		return extractorErrorf("artifact codec must not be nil")
	}
	build := func(request pipeline.BuildRequest) (*Extractor, error) {
		options, err := DecodeOptions(request.Options)
		if err != nil {
			return nil, err
		}
		return New(request.Dependencies.LLM, options)
	}
	return pipeline.RegisterExtractorBuilderWithRawAdapter(registry, ModuleSpec(), validateOptions,
		func(request pipeline.BuildRequest) (contracts.Extractor[dnd.SpellList], error) {
			return build(request)
		},
		func(request pipeline.BuildRequest) (contracts.LegacyRawExtractor, error) {
			extractor, err := build(request)
			if err != nil {
				return nil, err
			}
			return &rawAdapter{extractor: extractor, codec: codec}, nil
		},
	)
}

func (adapter *rawAdapter) Key() string { return Key }

func (adapter *rawAdapter) ReferenceSlots() []contracts.ReferenceSlot {
	return adapter.extractor.ReferenceSlots()
}

func (adapter *rawAdapter) ManifestMetadata() map[string]any {
	return adapter.extractor.ManifestMetadata()
}

func (adapter *rawAdapter) Extract(ctx context.Context, req contracts.ExtractionRequest) (contracts.ExtractionResult, error) {
	result, err := adapter.extractor.Extract(ctx, contracts.TypedExtractionRequest{
		Source: req.Source, Chunk: req.Chunk, AmbientContext: req.AmbientContext,
		SourceInput: req.SourceInput, SessionID: req.SessionID, References: req.References,
		LLMProfile: req.LLMProfile, Metadata: req.Metadata,
	})
	if err != nil {
		return contracts.ExtractionResult{}, err
	}
	content, err := adapter.codec.EncodeCandidate(result.Value)
	if err != nil {
		return contracts.ExtractionResult{}, extractorErrorf("encode canonical output: %w", err)
	}
	schema := adapter.codec.Schema()
	return contracts.ExtractionResult{
		Output: contracts.ExtractOutput{
			Schema:  contracts.ResponseSchema{ID: schema.ID, Name: schema.Name, Version: schema.Version, JSONSchema: append([]byte(nil), schema.JSONSchema...)},
			Payload: contracts.RawPayload{Content: content, MediaType: adapter.codec.MediaType(), Metadata: map[string]any{"spell_cast_count": len(result.Value.SpellCasts)}},
		},
		Warnings: result.Warnings,
	}, nil
}

func validateOptions(options map[string]any) error {
	_, err := DecodeOptions(options)
	return err
}

func DecodeOptions(options map[string]any) (Options, error) {
	if err := pipeline.RejectUnknownOptions(options); err != nil {
		return Options{}, extractorErrorf("%w", err)
	}
	return Options{}, nil
}

func extractorErrorf(format string, args ...any) error {
	return fmt.Errorf("dnd spells extractor: "+format, args...)
}