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"`
	Debug            *LLMDebugMaterial `json:"debug,omitempty"`
}

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

type LLMDebugMaterial struct {
	Prompt   *LLMDebugPrompt   `json:"prompt,omitempty"`
	Response *LLMDebugResponse `json:"response,omitempty"`
}

type LLMDebugPrompt struct {
	PromptID             string            `json:"prompt_id,omitempty"`
	PromptVersion        string            `json:"prompt_version,omitempty"`
	PromptHash           string            `json:"prompt_hash,omitempty"`
	SelectedProfileID    string            `json:"selected_profile_id,omitempty"`
	SelectedBackendID    string            `json:"selected_backend_id,omitempty"`
	SessionID            string            `json:"session_id,omitempty"`
	RenderedPromptHash   string            `json:"rendered_prompt_hash,omitempty"`
	Messages             []LLMDebugMessage `json:"messages,omitempty"`
	EffectiveModelParams map[string]any    `json:"effective_model_params,omitempty"`
	OutputContract       map[string]any    `json:"output_contract,omitempty"`
	StructuredOutput     map[string]any    `json:"structured_output,omitempty"`
	InputHashes          map[string]string `json:"input_hashes,omitempty"`
}

type LLMDebugMessage struct {
	Role         string         `json:"role"`
	Content      string         `json:"content"`
	CacheControl map[string]any `json:"cache_control,omitempty"`
}

type LLMDebugResponse struct {
	Content              string            `json:"content,omitempty"`
	RunID                string            `json:"run_id,omitempty"`
	PromptID             string            `json:"prompt_id,omitempty"`
	PromptVersion        string            `json:"prompt_version,omitempty"`
	PromptHash           string            `json:"prompt_hash,omitempty"`
	RenderedPromptHash   string            `json:"rendered_prompt_hash,omitempty"`
	SelectedProfileID    string            `json:"selected_profile_id,omitempty"`
	ModelName            string            `json:"model_name,omitempty"`
	Endpoint             string            `json:"endpoint,omitempty"`
	EffectiveModelParams map[string]any    `json:"effective_model_params,omitempty"`
	InputHashes          map[string]string `json:"input_hashes,omitempty"`
	Validation           map[string]any    `json:"validation,omitempty"`
	Usage                LLMDebugUsage     `json:"usage,omitempty"`
}

type LLMDebugUsage struct {
	PromptTokens     int `json:"prompt_tokens,omitempty"`
	CompletionTokens int `json:"completion_tokens,omitempty"`
	TotalTokens      int `json:"total_tokens,omitempty"`
	CachedTokens     int `json:"cached_tokens,omitempty"`
	CacheWriteTokens int `json:"cache_write_tokens,omitempty"`
}

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"`
	Metadata   map[string]any `json:"metadata,omitempty"`
}

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

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

type ChunkPlanResult struct {
	Plan     source.ChunkPlan `json:"plan"`
	Warnings []Warning        `json:"warnings,omitempty"`
}

type Chunker interface {
	Key() string
	ReferenceSlots() []ReferenceSlot
	Plan(ctx context.Context, req ChunkRequest) (ChunkPlanResult, 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"`
	AcceptedArtifactKinds []ArtifactKind `json:"accepted_artifact_kinds,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...)
		slot.AcceptedArtifactKinds = append([]ArtifactKind(nil), slot.AcceptedArtifactKinds...)
		out[i] = slot
	}
	return out
}

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

// ReferenceProducer identifies the operation that produced a generated
// reference. It contains provenance only; referenced bytes remain in Content.
type ReferenceProducer struct {
	PipelineID string `json:"pipeline_id,omitempty"`
	StepID     string `json:"step_id,omitempty"`
	LaneID     string `json:"lane_id,omitempty"`
	ModuleKey  string `json:"module_key,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"`
	ArtifactKind   ArtifactKind      `json:"artifact_kind,omitempty"`
	ArtifactSchema ArtifactSchema    `json:"artifact_schema,omitempty"`
	Producer       ReferenceProducer `json:"producer,omitempty"`
}

func CloneReferenceItem(item ReferenceItem) ReferenceItem {
	item.Content = append([]byte(nil), item.Content...)
	item.ArtifactSchema = CloneArtifactSchema(item.ArtifactSchema)
	return item
}

func (item ReferenceItem) MarshalJSON() ([]byte, error) {
	type referenceItemJSON struct {
		SlotName       string             `json:"slot_name"`
		MediaType      string             `json:"media_type,omitempty"`
		Digest         string             `json:"digest,omitempty"`
		Origin         ReferenceOrigin    `json:"origin"`
		SizeBytes      int64              `json:"size_bytes,omitempty"`
		BindingSource  string             `json:"binding_source,omitempty"`
		ArtifactKind   ArtifactKind       `json:"artifact_kind,omitempty"`
		ArtifactSchema *ArtifactSchema    `json:"artifact_schema,omitempty"`
		Producer       *ReferenceProducer `json:"producer,omitempty"`
	}
	encoded := referenceItemJSON{
		SlotName:      item.SlotName,
		MediaType:     item.MediaType,
		Digest:        item.Digest,
		Origin:        item.Origin,
		SizeBytes:     item.SizeBytes,
		BindingSource: item.BindingSource,
		ArtifactKind:  item.ArtifactKind,
	}
	if item.ArtifactSchema.ID != "" || item.ArtifactSchema.Name != "" || item.ArtifactSchema.Version != "" || len(item.ArtifactSchema.JSONSchema) > 0 {
		schema := CloneArtifactSchema(item.ArtifactSchema)
		encoded.ArtifactSchema = &schema
	}
	if item.Producer.PipelineID != "" || item.Producer.StepID != "" || item.Producer.LaneID != "" || item.Producer.ModuleKey != "" {
		producer := item.Producer
		encoded.Producer = &producer
	}
	return json.Marshal(encoded)
}

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

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

type ExecutionClass string

const (
	ExecutionClassDeterministic ExecutionClass = "deterministic"
	ExecutionClassLLMBacked     ExecutionClass = "llm_backed"
)

type ValidationResult struct {
	Approved               bool      `json:"approved"`
	ReasonCode             string    `json:"reason_code,omitempty"`
	Message                string    `json:"message,omitempty"`
	DiagnosticArtifactPath string    `json:"diagnostic_artifact_path,omitempty"`
	Warnings               []Warning `json:"warnings,omitempty"`
}

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 []SerializedOutput    `json:"normalize_outputs,omitempty"`
	Rejected         []RejectedOutput      `json:"rejected,omitempty"`
	Warnings         []Warning             `json:"warnings,omitempty"`
	LLMProfile       string                `json:"llm_profile,omitempty"`
	Metadata         map[string]any        `json:"metadata,omitempty"`
	ChunkMap         *SerializedArtifact   `json:"chunk_map,omitempty"`
	EvidenceContext  *SerializedArtifact   `json:"evidence_context,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"`
	StepID                 string `json:"step_id,omitempty"`
	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
}