package main

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

	"github.com/spf13/cobra"

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

var (
	queryTarget  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().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 target.
	targetName := queryTarget
	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 client (use first endpoint directly — no router needed for a one-shot query).
	ep := target.Endpoints[0]
	apiKey := ep.APIKey
	if apiKey == "" {
		apiKey = target.APIKey
	}
	timeout := time.Duration(target.TimeoutSecs) * time.Second
	if timeout == 0 {
		timeout = 30 * time.Second
	}
	httpClient := &http.Client{Timeout: timeout}

	var client embedclient.Client
	switch target.APIType {
	case "google":
		client = embedclient.NewGoogle(ep.URL, apiKey, target.Model, httpClient)
	default:
		client = embedclient.NewOpenAI(ep.URL, apiKey, httpClient)
	}

	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()

	resp, err := client.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 {
		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
}
