package main

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"time"

	"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/embedclient"
)

var (
	queryTarget  string
	queryMap     string
	queryRaw     bool
	queryCompact bool
)

var queryCmd = &cobra.Command{
	Use:   "query <text>",
	Short: "Embed text and print the resulting vector",
	Long: `Sends text to the configured forwarding target, applies the dimension adapter,
and prints the resulting vector as a JSON array.

Text can be supplied as a positional argument or via stdin (use - as the argument).`,
	Args: cobra.MaximumNArgs(1),
	RunE: runQuery,
}

func init() {
	queryCmd.Flags().StringVar(&queryTarget, "target", "",
		"forward target to use (default: forward.default from config)")
	queryCmd.Flags().StringVar(&queryMap, "map", "",
		"extra_map key to use; overrides both adapter and forward target")
	queryCmd.Flags().BoolVar(&queryRaw, "raw", false,
		"skip the adapter — output the raw vector from the backing model")
	queryCmd.Flags().BoolVar(&queryCompact, "compact", false,
		"print vector on a single line instead of pretty-printed")
	rootCmd.AddCommand(queryCmd)
}

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

	// Resolve text: positional arg, "-" reads stdin, no arg reads stdin.
	text, err := queryText(args)
	if err != nil {
		return err
	}

	// Resolve adapter and forward target — --map overrides both.
	var adp adapter.Adapter
	targetName := queryTarget
	if queryMap != "" {
		mc, ok := cfg.ExtraMaps[queryMap]
		if !ok {
			return fmt.Errorf("extra_map %q not found in config", queryMap)
		}
		extraMaps, buildErr := buildExtraMapAdapters(cfg)
		if buildErr != nil {
			return fmt.Errorf("build extra_maps: %w", buildErr)
		}
		adp = extraMaps[queryMap].Adapter
		if mc.ForwardTarget != "" && targetName == "" {
			targetName = mc.ForwardTarget
		}
	}

	if targetName == "" {
		targetName = cfg.Forward.Default
	}
	target, ok := cfg.Forward.Targets[targetName]
	if !ok {
		return fmt.Errorf("target %q not found in config", targetName)
	}
	if len(target.Endpoints) == 0 {
		return fmt.Errorf("target %q has no endpoints", targetName)
	}

	// Build a full router so retry/demerit logic applies identically to the server path.
	client, err := buildClients(cfg)
	if err != nil {
		return fmt.Errorf("build clients: %w", err)
	}
	embedClient, ok := client[targetName]
	if !ok {
		return fmt.Errorf("target %q not found after building clients", targetName)
	}

	timeout := time.Duration(target.TimeoutSecs) * time.Second
	if timeout == 0 {
		timeout = 30 * time.Second
	}
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	resp, err := embedClient.Embed(ctx, embedclient.Request{
		Texts: []string{text},
		Model: target.Model,
	})
	if err != nil {
		return fmt.Errorf("embed: %w", err)
	}
	if len(resp.Embeddings) == 0 || len(resp.Embeddings[0]) == 0 {
		return fmt.Errorf("empty embedding in response")
	}

	vec := resp.Embeddings[0]

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

	fmt.Fprintf(os.Stderr, "target=%s  model=%s  dims=%d  tokens=%d\n",
		targetName, resp.Model, len(vec), resp.Usage.TotalTokens)

	return printVector(vec)
}

func queryText(args []string) (string, error) {
	if len(args) == 0 || args[0] == "-" {
		raw, err := os.ReadFile("/dev/stdin")
		if err != nil {
			return "", fmt.Errorf("read stdin: %w", err)
		}
		return strings.TrimRight(string(raw), "\n"), nil
	}
	return args[0], nil
}

func printVector(vec []float32) error {
	// Convert to []any so json.Marshal produces clean floats without float32 quirks.
	out := make([]float64, len(vec))
	for i, v := range vec {
		out[i] = float64(v)
	}

	var b []byte
	var err error
	if queryCompact {
		b, err = json.Marshal(out)
	} else {
		b, err = json.MarshalIndent(out, "", "  ")
	}
	if err != nil {
		return fmt.Errorf("marshal vector: %w", err)
	}

	fmt.Println(string(b))
	return nil
}
