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"`
	StructuredOutputRepairAttempts *int                `json:"structured_output_repair_attempts,omitempty"`
	Correction                     *SemanticCorrection `json:"-"`
}

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"`
	RepairAttempts   int               `json:"repair_attempts,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"`
	ProviderError        *LLMDebugProviderError `json:"provider_error,omitempty"`
}

type LLMDebugProviderError struct {
	StatusCode int    `json:"status_code,omitempty"`
	Code       string `json:"code,omitempty"`
	Type       string `json:"type,omitempty"`
	Message    string `json:"message,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"`
	StructuredOutputRepairAttempts *int           `json:"structured_output_repair_attempts,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"`
	StructuredOutputRepairAttempts *int                   `json:"structured_output_repair_attempts,omitempty"`
	Correction                     *SemanticCorrection    `json:"-"`
	Metadata                       map[string]any         `json:"metadata,omitempty"`
}

type ChunkPlanResult struct {
	Plan           source.ChunkPlan     `json:"plan"`
	Diagnostics    []ProducerDiagnostic `json:"diagnostics,omitempty"`
	ModelCandidate *ModelCandidate      `json:"-"`
}

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"`
	CorrectionGuidance     string               `json:"-"`
	DiagnosticArtifactPath string               `json:"diagnostic_artifact_path,omitempty"`
	Diagnostics            []ProducerDiagnostic `json:"diagnostics,omitempty"`
}

type OutputRequest struct {
	Manifest                       artifacts.RunManifest `json:"manifest"`
	NormalizeOutputs               []SerializedOutput    `json:"normalize_outputs,omitempty"`
	Rejected                       []RejectedOutput      `json:"rejected,omitempty"`
	Diagnostics                    DiagnosticCollection  `json:"diagnostics,omitempty"`
	LLMProfile                     string                `json:"llm_profile,omitempty"`
	StructuredOutputRepairAttempts *int                  `json:"structured_output_repair_attempts,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"`
}

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"`
	Validation             *artifacts.ValidationSummary `json:"validation,omitempty"`
}

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