package main

import (
	"encoding/json"
	"fmt"
	"io"
	"os"

	"github.com/spf13/cobra"

	"github.com/Warky-Devs/vecna.git/pkg/adapter"
	"github.com/Warky-Devs/vecna.git/pkg/config"
	"github.com/Warky-Devs/vecna.git/pkg/server"
)

var (
	convertInput  string
	convertOutput string
)

var convertCmd = &cobra.Command{
	Use:   "convert",
	Short: "Convert vectors from one dimension to another",
	Long:  "Reads a JSON array of float32 vectors, applies the configured adapter, and writes the result.",
	RunE:  runConvert,
}

func init() {
	convertCmd.Flags().StringVarP(&convertInput, "input", "i", "-", "input file path (- for stdin)")
	convertCmd.Flags().StringVarP(&convertOutput, "output", "o", "-", "output file path (- for stdout)")
}

func runConvert(cmd *cobra.Command, _ []string) error {
	cfg, err := config.Load(cfgFile)
	if err != nil {
		return fmt.Errorf("load config: %w", err)
	}

	adp, err := buildAdapter(cfg)
	if err != nil {
		return fmt.Errorf("build adapter: %w", err)
	}

	in, err := openReader(convertInput)
	if err != nil {
		return fmt.Errorf("open input: %w", err)
	}
	if f, ok := in.(*os.File); ok && f != os.Stdin {
		defer func() { _ = f.Close() }()
	}

	out, err := openWriter(convertOutput)
	if err != nil {
		return fmt.Errorf("open output: %w", err)
	}
	if f, ok := out.(*os.File); ok && f != os.Stdout {
		defer func() { _ = f.Close() }()
	}

	var vecs [][]float32
	if err := json.NewDecoder(in).Decode(&vecs); err != nil {
		return fmt.Errorf("decode input: %w", err)
	}

	result := make([][]float32, len(vecs))
	for i, v := range vecs {
		adapted, adaptErr := adp.Adapt(v)
		if adaptErr != nil {
			return fmt.Errorf("adapt vector %d: %w", i, adaptErr)
		}
		result[i] = adapted
	}

	enc := json.NewEncoder(out)
	enc.SetIndent("", "  ")
	if err := enc.Encode(result); err != nil {
		return fmt.Errorf("encode output: %w", err)
	}
	return nil
}

func openReader(path string) (io.Reader, error) {
	if path == "-" {
		return os.Stdin, nil
	}
	return os.Open(path)
}

func openWriter(path string) (io.Writer, error) {
	if path == "-" {
		return os.Stdout, nil
	}
	return os.Create(path)
}

// buildAdapter constructs the Adapter from the loaded config.
func buildAdapter(cfg *config.Config) (adapter.Adapter, error) {
	return buildAdapterFromConfig(cfg.Adapter)
}

// buildExtraMapAdapters builds a server.ExtraMap for each entry in cfg.ExtraMaps.
// Unset adapter fields fall back to the global Adapter values.
func buildExtraMapAdapters(cfg *config.Config) (map[string]server.ExtraMap, error) {
	result := make(map[string]server.ExtraMap, len(cfg.ExtraMaps))
	for name, mc := range cfg.ExtraMaps {
		ac := config.AdapterConfig{
			Type:         coalesce(mc.Type, cfg.Adapter.Type),
			SourceDim:    coalescei(mc.SourceDim, cfg.Adapter.SourceDim),
			TargetDim:    mc.TargetDim,
			TruncateMode: coalesce(mc.TruncateMode, cfg.Adapter.TruncateMode),
			PadMode:      coalesce(mc.PadMode, cfg.Adapter.PadMode),
			Seed:         coalescei64(mc.Seed, cfg.Adapter.Seed),
			MatrixFile:   coalesce(mc.MatrixFile, cfg.Adapter.MatrixFile),
		}
		adp, err := buildAdapterFromConfig(ac)
		if err != nil {
			return nil, fmt.Errorf("extra_map %q: %w", name, err)
		}
		result[name] = server.ExtraMap{
			Adapter:       adp,
			ForwardTarget: mc.ForwardTarget,
		}
	}
	return result, nil
}

// buildAdapterFromConfig constructs an Adapter from an AdapterConfig.
func buildAdapterFromConfig(ac config.AdapterConfig) (adapter.Adapter, error) {
	if ac.SourceDim == 0 && ac.TargetDim == 0 {
		return adapter.NewPassthrough(), nil
	}
	switch ac.Type {
	case "truncate":
		tm, pm, err := parseTruncateModes(ac.TruncateMode, ac.PadMode)
		if err != nil {
			return nil, err
		}
		return adapter.NewTruncate(ac.SourceDim, ac.TargetDim, tm, pm)

	case "random":
		return adapter.NewRandom(ac.SourceDim, ac.TargetDim, ac.Seed)

	case "projection":
		if ac.MatrixFile == "" {
			return nil, fmt.Errorf("adapter type 'projection' requires matrix_file")
		}
		matrix, err := loadMatrix(ac.MatrixFile)
		if err != nil {
			return nil, fmt.Errorf("load projection matrix: %w", err)
		}
		return adapter.NewProjection(ac.SourceDim, ac.TargetDim, matrix)

	default:
		return nil, fmt.Errorf("unknown adapter type %q; valid: truncate, random, projection", ac.Type)
	}
}

func coalesce(a, b string) string {
	if a != "" {
		return a
	}
	return b
}

func coalescei(a, b int) int {
	if a != 0 {
		return a
	}
	return b
}

func coalescei64(a, b int64) int64 {
	if a != 0 {
		return a
	}
	return b
}

func parseTruncateModes(truncMode, padMode string) (adapter.TruncateMode, adapter.PadMode, error) {
	var tm adapter.TruncateMode
	switch truncMode {
	case "from_end", "":
		tm = adapter.TruncateFromEnd
	case "from_start":
		tm = adapter.TruncateFromStart
	default:
		return 0, 0, fmt.Errorf("unknown truncate_mode %q; valid: from_end, from_start", truncMode)
	}

	var pm adapter.PadMode
	switch padMode {
	case "at_end", "":
		pm = adapter.PadAtEnd
	case "at_start":
		pm = adapter.PadAtStart
	default:
		return 0, 0, fmt.Errorf("unknown pad_mode %q; valid: at_end, at_start", padMode)
	}

	return tm, pm, nil
}

func loadMatrix(path string) ([][]float32, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer func() { _ = f.Close() }()

	var m [][]float32
	if err := json.NewDecoder(f).Decode(&m); err != nil {
		return nil, fmt.Errorf("decode matrix JSON: %w", err)
	}
	return m, nil
}
