package config

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

const SupportedFileConfigVersion = 4

type Config struct {
	PromptKit   PromptKitConfig                     `json:"promptkit,omitempty"`
	Pipelines   map[string]pipeline.PipelineProfile `json:"pipelines"`
	Concurrency ConcurrencyConfig                   `json:"concurrency"`
	Output      OutputConfig                        `json:"output"`
	Cache       CacheConfig                         `json:"cache"`
	Debug       DebugConfig                         `json:"debug"`
}

type PromptKitConfig struct {
	ProfileDir   string                       `json:"profile_dir,omitempty"`
	ProfileFile  string                       `json:"profile_file,omitempty"`
	LocalBackend *PromptKitLocalBackendConfig `json:"local_backend,omitempty"`
}

type PromptKitLocalBackendConfig struct {
	Endpoint         string `json:"endpoint"`
	ConcurrencyLimit int    `json:"concurrency_limit"`
}

type ConcurrencyConfig struct {
	TotalLLM     int            `json:"total_llm"`
	StageWorkers map[string]int `json:"stage_workers"`

	extractWorkersConfigured bool
	defaultedExtractWorkers  int
}

const defaultLLMConcurrency = 16

type OutputConfig struct {
	Directory string `json:"directory"`
}

type CacheConfig struct {
	ChunkPlans  ChunkPlanCacheConfig  `json:"chunk_plans"`
	Checkpoints CheckpointCacheConfig `json:"checkpoints"`
}

type ChunkPlanCacheConfig struct {
	Directory string                  `json:"directory,omitempty"`
	Mode      pipeline.ChunkCacheMode `json:"mode"`
}

type CheckpointCacheConfig struct {
	Enabled   bool   `json:"enabled"`
	Directory string `json:"directory,omitempty"`
}
type DebugConfig struct {
	Directory string `json:"directory"`
}

func Default() Config {
	return Config{
		Pipelines: map[string]pipeline.PipelineProfile{},
		Concurrency: ConcurrencyConfig{
			TotalLLM:                defaultLLMConcurrency,
			StageWorkers:            map[string]int{"extract": defaultLLMConcurrency},
			defaultedExtractWorkers: defaultLLMConcurrency,
		},
		Output: OutputConfig{Directory: "./notarius-output"},
		Cache:  CacheConfig{ChunkPlans: ChunkPlanCacheConfig{Mode: pipeline.ChunkCacheAuto}},
		Debug:  DebugConfig{Directory: "./notarius-debug"},
	}
}

func cloneConfig(in Config) Config {
	out := in
	if in.PromptKit.LocalBackend != nil {
		localBackend := *in.PromptKit.LocalBackend
		out.PromptKit.LocalBackend = &localBackend
	}
	out.Concurrency.StageWorkers = cloneIntMap(in.Concurrency.StageWorkers)
	out.Pipelines = make(map[string]pipeline.PipelineProfile, len(in.Pipelines))
	for key, profile := range in.Pipelines {
		out.Pipelines[key] = clonePipelineProfile(profile)
	}
	return out
}

func cloneIntMap(in map[string]int) map[string]int {
	if len(in) == 0 {
		return nil
	}
	out := make(map[string]int, len(in))
	for key, value := range in {
		out[key] = value
	}
	return out
}

func (c *ConcurrencyConfig) recomputeStageWorkerDefaults() {
	if c == nil {
		return
	}
	if c.StageWorkers == nil {
		c.StageWorkers = make(map[string]int)
	}
	if !c.extractWorkersConfigured {
		if value, ok := c.StageWorkers["extract"]; ok && (c.defaultedExtractWorkers == 0 || value != c.defaultedExtractWorkers) {
			c.extractWorkersConfigured = true
			return
		}
		c.StageWorkers["extract"] = c.TotalLLM
		c.defaultedExtractWorkers = c.TotalLLM
	}
}

func clonePipelineProfile(in pipeline.PipelineProfile) pipeline.PipelineProfile {
	out := in
	out.ValidationPolicy = cloneValidationPolicyOverride(in.ValidationPolicy)
	if in.StructuredOutputRepairAttempts != nil {
		value := *in.StructuredOutputRepairAttempts
		out.StructuredOutputRepairAttempts = &value
	}
	out.Input = cloneModuleBinding(in.Input)
	out.Chunk = cloneModuleBinding(in.Chunk)
	out.Output = cloneModuleBinding(in.Output)
	out.References = cloneReferenceSourceMap(in.References)
	if len(in.Artifacts) > 0 {
		out.Artifacts = make(map[string]pipeline.ArtifactLaneProfile, len(in.Artifacts))
		for key, lane := range in.Artifacts {
			out.Artifacts[key] = cloneArtifactLaneProfile(lane)
		}
	}
	if in.Steps != nil {
		out.Steps = make([]pipeline.PipelineStepProfile, len(in.Steps))
		for i, step := range in.Steps {
			out.Steps[i] = clonePipelineStepProfile(step)
		}
	}
	return out
}

func clonePipelineStepProfile(in pipeline.PipelineStepProfile) pipeline.PipelineStepProfile {
	out := in
	out.ID = in.ID
	out.References = cloneReferenceSourceMap(in.References)
	if len(in.Artifacts) > 0 {
		out.Artifacts = make(map[string]pipeline.ArtifactLaneProfile, len(in.Artifacts))
		for key, lane := range in.Artifacts {
			out.Artifacts[key] = cloneArtifactLaneProfile(lane)
		}
	}
	return out
}

func cloneArtifactLaneProfile(in pipeline.ArtifactLaneProfile) pipeline.ArtifactLaneProfile {
	out := in
	out.Extract = cloneModuleBinding(in.Extract)
	out.Merge = cloneModuleBinding(in.Merge)
	out.Normalize = cloneModuleBinding(in.Normalize)
	out.References = cloneReferenceSourceMap(in.References)
	if len(in.Validators) > 0 {
		out.Validators = make([]pipeline.ModuleBinding, len(in.Validators))
		for i, binding := range in.Validators {
			out.Validators[i] = cloneModuleBinding(binding)
		}
	}
	return out
}

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

func cloneReferenceSourceMap(in map[string]pipeline.ReferenceSource) map[string]pipeline.ReferenceSource {
	if len(in) == 0 {
		return nil
	}
	out := make(map[string]pipeline.ReferenceSource, len(in))
	for key, source := range in {
		out[key] = cloneReferenceSource(source)
	}
	return out
}

func cloneReferenceSource(in pipeline.ReferenceSource) pipeline.ReferenceSource {
	out := in
	if in.Artifact != nil {
		artifact := *in.Artifact
		out.Artifact = &artifact
	}
	return out
}

func cloneModuleBinding(in pipeline.ModuleBinding) pipeline.ModuleBinding {
	out := in
	out.ValidationPolicy = cloneValidationPolicyOverride(in.ValidationPolicy)
	if in.StructuredOutputRepairAttempts != nil {
		value := *in.StructuredOutputRepairAttempts
		out.StructuredOutputRepairAttempts = &value
	}
	if len(in.Options) > 0 {
		out.Options = cloneOptions(in.Options)
	}
	out.References = cloneReferenceSourceMap(in.References)
	out.Validators = cloneValidatorOverride(in.Validators)
	return out
}

func cloneValidationPolicyOverride(in *pipeline.ValidationPolicyOverride) *pipeline.ValidationPolicyOverride {
	if in == nil {
		return nil
	}
	out := *in
	if in.ProducerStructuralFailure != nil {
		value := *in.ProducerStructuralFailure
		out.ProducerStructuralFailure = &value
	}
	if in.SemanticRejection != nil {
		value := *in.SemanticRejection
		out.SemanticRejection = &value
	}
	if in.ValidatorFailure != nil {
		value := *in.ValidatorFailure
		out.ValidatorFailure = &value
	}
	return &out
}

func cloneValidatorOverride(in pipeline.ValidatorOverride) pipeline.ValidatorOverride {
	out := pipeline.ValidatorOverride{Set: in.Set}
	if len(in.Validators) > 0 {
		out.Validators = make([]pipeline.ModuleBinding, len(in.Validators))
		for i, binding := range in.Validators {
			out.Validators[i] = cloneModuleBinding(binding)
		}
	}
	return out
}

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

func cloneOptionValue(value any) any {
	switch typed := value.(type) {
	case map[string]any:
		return cloneOptions(typed)
	case []any:
		out := make([]any, len(typed))
		for i, item := range typed {
			out[i] = cloneOptionValue(item)
		}
		return out
	default:
		return typed
	}
}
