package contracts

import (
	"context"
	"encoding/json"

	"gitea.maximumdirect.net/eric/notarius/internal/core/artifacts"
	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
)

type StructuredCompletionRequest struct {
	StageName     string         `json:"stage_name"`
	PromptID      string         `json:"prompt_id,omitempty"`
	PromptVersion string         `json:"prompt_version,omitempty"`
	ProfileID     string         `json:"profile_id,omitempty"`
	SessionID     string         `json:"session_id,omitempty"`
	Inputs        LLMInputSet    `json:"inputs,omitempty"`
	Vars          map[string]any `json:"vars,omitempty"`
}

type StructuredCompletionResponse struct {
	Content          json.RawMessage `json:"content"`
	Provider         string          `json:"provider,omitempty"`
	Model            string          `json:"model,omitempty"`
	ProfileID        string          `json:"profile_id,omitempty"`
	PromptTokens     int             `json:"prompt_tokens,omitempty"`
	CompletionTokens int             `json:"completion_tokens,omitempty"`
	TotalTokens      int             `json:"total_tokens,omitempty"`
}

type StructuredLLMClient interface {
	CompleteStructured(ctx context.Context, req StructuredCompletionRequest, out any) (StructuredCompletionResponse, error)
}

type LLMProfileManifestProvider interface {
	LLMProfileManifests() []artifacts.LLMProfileManifest
}

type LLMInputMaterial struct {
	Name      string `json:"name"`
	MediaType string `json:"media_type,omitempty"`
	Content   []byte `json:"-"`
	Digest    string `json:"digest,omitempty"`
	OriginURI string `json:"origin_uri,omitempty"`
	SizeBytes int64  `json:"size_bytes,omitempty"`
}

func NewLLMInputMaterial(name string, mediaType string, content []byte, digest string, originURI string) LLMInputMaterial {
	return LLMInputMaterial{
		Name:      name,
		MediaType: mediaType,
		Content:   append([]byte(nil), content...),
		Digest:    digest,
		OriginURI: originURI,
		SizeBytes: int64(len(content)),
	}
}

func (material LLMInputMaterial) Clone() LLMInputMaterial {
	material.Content = append([]byte(nil), material.Content...)
	return material
}

type LLMInputSet map[string]LLMInputMaterial

func (set LLMInputSet) Clone() LLMInputSet {
	if len(set) == 0 {
		return nil
	}
	out := make(LLMInputSet, len(set))
	for key, material := range set {
		out[key] = material.Clone()
	}
	return out
}

type ParseRequest struct {
	SourceID   string         `json:"source_id,omitempty"`
	Path       string         `json:"path,omitempty"`
	Raw        []byte         `json:"-"`
	LLMProfile string         `json:"llm_profile,omitempty"`
	Options    map[string]any `json:"options,omitempty"`
	Metadata   map[string]any `json:"metadata,omitempty"`
}

type InputAdapter interface {
	Key() string
	Parse(ctx context.Context, req ParseRequest) (*source.SourceDocument, error)
}

type SourceChunk struct {
	ID          string              `json:"id"`
	SourceID    string              `json:"source_id"`
	Index       int                 `json:"index"`
	StartUnitID int                 `json:"start_unit_id"`
	EndUnitID   int                 `json:"end_unit_id"`
	Content     []byte              `json:"-"`
	MediaType   string              `json:"media_type"`
	Units       []source.SourceUnit `json:"units"`
	Metadata    map[string]any      `json:"metadata,omitempty"`
}

type ChunkRequest struct {
	Source      *source.SourceDocument `json:"-"`
	SourceInput LLMInputMaterial       `json:"source_input,omitempty"`
	SessionID   string                 `json:"session_id,omitempty"`
	References  ReferenceSet           `json:"references,omitempty"`
	LLMClient   StructuredLLMClient    `json:"-"`
	LLMProfile  string                 `json:"llm_profile,omitempty"`
	Options     map[string]any         `json:"options,omitempty"`
	Metadata    map[string]any         `json:"metadata,omitempty"`
}

type ChunkResult struct {
	Chunks   []SourceChunk `json:"chunks"`
	Warnings []Warning     `json:"warnings,omitempty"`
}

type Chunker interface {
	Key() string
	ReferenceSlots() []ReferenceSlot
	Chunk(ctx context.Context, req ChunkRequest) (ChunkResult, error)
}

const (
	ReferenceBindingSourceConfig = "config"
	ReferenceBindingSourceCLI    = "cli"
)

type ReferenceSlot struct {
	Name               string   `json:"name"`
	Description        string   `json:"description,omitempty"`
	Required           bool     `json:"required,omitempty"`
	AcceptedMediaTypes []string `json:"accepted_media_types,omitempty"`
	Multiple           bool     `json:"multiple,omitempty"`
	MaxBytes           int64    `json:"max_bytes,omitempty"`
}

func CloneReferenceSlots(slots []ReferenceSlot) []ReferenceSlot {
	if len(slots) == 0 {
		return nil
	}
	out := make([]ReferenceSlot, len(slots))
	for i, slot := range slots {
		slot.AcceptedMediaTypes = append([]string(nil), slot.AcceptedMediaTypes...)
		out[i] = slot
	}
	return out
}

type ReferenceOrigin struct {
	Type string `json:"type"`
	URI  string `json:"uri,omitempty"`
}

type ReferenceItem struct {
	SlotName      string          `json:"slot_name"`
	MediaType     string          `json:"media_type,omitempty"`
	Content       []byte          `json:"-"`
	Digest        string          `json:"digest,omitempty"`
	Origin        ReferenceOrigin `json:"origin"`
	SizeBytes     int64           `json:"size_bytes,omitempty"`
	BindingSource string          `json:"binding_source,omitempty"`
}

type ResolvedReferenceSlot struct {
	Slot  ReferenceSlot   `json:"slot"`
	Items []ReferenceItem `json:"items,omitempty"`
}

type ReferenceSet struct {
	Slots map[string]ResolvedReferenceSlot `json:"slots,omitempty"`
}

type ExtractionRequest struct {
	Source         *source.SourceDocument `json:"-"`
	Chunk          *SourceChunk           `json:"chunk,omitempty"`
	AmbientContext map[string]any         `json:"ambient_context,omitempty"`
	SourceInput    LLMInputMaterial       `json:"source_input,omitempty"`
	SessionID      string                 `json:"session_id,omitempty"`
	References     ReferenceSet           `json:"references,omitempty"`
	LLMClient      StructuredLLMClient    `json:"-"`
	LLMProfile     string                 `json:"llm_profile,omitempty"`
	Options        map[string]any         `json:"options,omitempty"`
	Metadata       map[string]any         `json:"metadata,omitempty"`
}

type ExtractionResult struct {
	Output   ExtractOutput `json:"output"`
	Warnings []Warning     `json:"warnings,omitempty"`
}

type Extractor interface {
	Key() string
	ReferenceSlots() []ReferenceSlot
	Extract(ctx context.Context, req ExtractionRequest) (ExtractionResult, error)
}

type RawPayload struct {
	Content   []byte         `json:"-"`
	MediaType string         `json:"media_type"`
	Metadata  map[string]any `json:"metadata,omitempty"`
	Warnings  []Warning      `json:"warnings,omitempty"`
}

type ResponseSchema struct {
	ID      string `json:"id,omitempty"`
	Name    string `json:"name,omitempty"`
	Version string `json:"version,omitempty"`
}

type ExtractOutput struct {
	LaneID       string         `json:"lane_id"`
	ExtractorKey string         `json:"extractor_key"`
	SourceID     string         `json:"source_id"`
	ChunkID      string         `json:"chunk_id"`
	ChunkIndex   int            `json:"chunk_index"`
	Schema       ResponseSchema `json:"schema,omitempty"`
	Payload      RawPayload     `json:"payload"`
}

type MergeRequest struct {
	Source         *source.SourceDocument `json:"-"`
	LaneID         string                 `json:"lane_id"`
	ExtractOutputs []ExtractOutput        `json:"extract_outputs"`
	SourceInput    LLMInputMaterial       `json:"source_input,omitempty"`
	SessionID      string                 `json:"session_id,omitempty"`
	References     ReferenceSet           `json:"references,omitempty"`
	LLMClient      StructuredLLMClient    `json:"-"`
	LLMProfile     string                 `json:"llm_profile,omitempty"`
	Options        map[string]any         `json:"options,omitempty"`
	Metadata       map[string]any         `json:"metadata,omitempty"`
}

type MergeResult struct {
	Output   MergeOutput `json:"output"`
	Warnings []Warning   `json:"warnings,omitempty"`
}

type MergeOutput struct {
	LaneID    string         `json:"lane_id"`
	MergerKey string         `json:"merger_key"`
	SourceID  string         `json:"source_id,omitempty"`
	Schema    ResponseSchema `json:"schema,omitempty"`
	Payload   RawPayload     `json:"payload"`
}

type Merger interface {
	Key() string
	Merge(ctx context.Context, req MergeRequest) (MergeResult, error)
}

type NormalizeRequest struct {
	Source      *source.SourceDocument `json:"-"`
	LaneID      string                 `json:"lane_id"`
	MergeOutput MergeOutput            `json:"merge_output"`
	SourceInput LLMInputMaterial       `json:"source_input,omitempty"`
	SessionID   string                 `json:"session_id,omitempty"`
	References  ReferenceSet           `json:"references,omitempty"`
	LLMClient   StructuredLLMClient    `json:"-"`
	LLMProfile  string                 `json:"llm_profile,omitempty"`
	Options     map[string]any         `json:"options,omitempty"`
	Metadata    map[string]any         `json:"metadata,omitempty"`
}

type NormalizeResult struct {
	Output   NormalizeOutput `json:"output"`
	Warnings []Warning       `json:"warnings,omitempty"`
}

type NormalizeOutput struct {
	LaneID        string         `json:"lane_id"`
	NormalizerKey string         `json:"normalizer_key"`
	SourceID      string         `json:"source_id,omitempty"`
	Schema        ResponseSchema `json:"schema,omitempty"`
	Payload       RawPayload     `json:"payload"`
}

type Normalizer interface {
	Key() string
	ReferenceSlots() []ReferenceSlot
	Normalize(ctx context.Context, req NormalizeRequest) (NormalizeResult, error)
}

type ValidationRequest struct {
	Source     *source.SourceDocument        `json:"-"`
	Candidates []artifacts.ArtifactCandidate `json:"candidates"`
	LLMProfile string                        `json:"llm_profile,omitempty"`
	Options    map[string]any                `json:"options,omitempty"`
	Metadata   map[string]any                `json:"metadata,omitempty"`
}

type ValidationDecision struct {
	CandidateIndex         int    `json:"candidate_index"`
	Approved               bool   `json:"approved"`
	ReasonCode             string `json:"reason_code"`
	Message                string `json:"message"`
	DiagnosticArtifactPath string `json:"diagnostic_artifact_path,omitempty"`
}

type ValidationResult struct {
	ValidatorName string               `json:"validator_name"`
	Decisions     []ValidationDecision `json:"decisions"`
	Warnings      []Warning            `json:"warnings,omitempty"`
}

type Validator interface {
	Name() string
	Validate(ctx context.Context, req ValidationRequest) (ValidationResult, error)
}

type Warning struct {
	Scope      string `json:"scope,omitempty"`
	ReasonCode string `json:"reason_code"`
	Message    string `json:"message"`
}

type OutputRequest struct {
	Manifest         artifacts.RunManifest `json:"manifest"`
	NormalizeOutputs []NormalizeOutput     `json:"normalize_outputs,omitempty"`
	Rejected         []RejectedOutput      `json:"rejected,omitempty"`
	Warnings         []Warning             `json:"warnings,omitempty"`
	LLMProfile       string                `json:"llm_profile,omitempty"`
	Options          map[string]any        `json:"options,omitempty"`
	Metadata         map[string]any        `json:"metadata,omitempty"`
}

type OutputFile struct {
	Name        string `json:"name"`
	ContentType string `json:"content_type,omitempty"`
	Bytes       []byte `json:"-"`
}

type OutputResult struct {
	Files    []OutputFile `json:"files,omitempty"`
	Warnings []Warning    `json:"warnings,omitempty"`
}

type OutputEncoder interface {
	Key() string
	Encode(ctx context.Context, req OutputRequest) (OutputResult, error)
}

type RejectedOutput struct {
	Stage                  string `json:"stage"`
	LaneID                 string `json:"lane_id,omitempty"`
	ModuleKey              string `json:"module_key,omitempty"`
	ChunkID                string `json:"chunk_id,omitempty"`
	ChunkIndex             int    `json:"chunk_index,omitempty"`
	ValidatorName          string `json:"validator_name,omitempty"`
	ReasonCode             string `json:"reason_code,omitempty"`
	Message                string `json:"message"`
	AttemptCount           int    `json:"attempt_count,omitempty"`
	DiagnosticArtifactPath string `json:"diagnostic_artifact_path,omitempty"`
}

type ManifestMetadataProvider interface {
	ManifestMetadata() map[string]any
}