package main

import (
	"context"
	"fmt"
	"math"
	"net/http"
	"time"

	"github.com/spf13/cobra"

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

var removeBroken bool

var testCmd = &cobra.Command{
	Use:   "test",
	Short: "Send a test embedding request to each configured forwarder",
	RunE:  runTest,
}

func init() {
	testCmd.Flags().BoolVar(&removeBroken, "remove-broken", false,
		"Remove endpoints (and targets with no endpoints left) that fail the test from the config file")
	rootCmd.AddCommand(testCmd)
}

const testPhrase = "The quick brown fox jumps over the lazy dog"

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

	if len(cfg.Forward.Targets) == 0 {
		fmt.Println("No forwarder targets configured.")
		return nil
	}

	fmt.Printf("Test phrase: %q\n\n", testPhrase)

	passed, failed := 0, 0
	// brokenEndpoints maps target name → list of failing endpoint URLs
	brokenEndpoints := make(map[string][]string)

	for targetName, target := range cfg.Forward.Targets {
		fmt.Printf("[ %s ]  model: %s  type: %s\n", targetName, target.Model, target.APIType)

		timeout := time.Duration(target.TimeoutSecs) * time.Second
		if timeout == 0 {
			timeout = 30 * time.Second
		}
		httpClient := &http.Client{Timeout: timeout}

		for _, ep := range target.Endpoints {
			apiKey := ep.APIKey
			if apiKey == "" {
				apiKey = target.APIKey
			}

			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)
			start := time.Now()
			resp, embedErr := client.Embed(ctx, embedclient.Request{
				Texts: []string{testPhrase},
				Model: target.Model,
			})
			elapsed := time.Since(start)
			cancel()

			if embedErr != nil {
				fmt.Printf("  %-45s  FAIL  %s\n", ep.URL, truncate(embedErr.Error(), 60))
				brokenEndpoints[targetName] = append(brokenEndpoints[targetName], ep.URL)
				failed++
				continue
			}

			dims, norm := embeddingStats(resp.Embeddings)
			fmt.Printf("  %-45s  OK    %dms  dims=%d  norm=%.4f\n",
				ep.URL, elapsed.Milliseconds(), dims, norm)
			passed++
		}

		if len(target.Endpoints) == 0 {
			fmt.Println("  (no endpoints configured)")
		}
		fmt.Println()
	}

	fmt.Printf("Results: %d passed, %d failed\n", passed, failed)

	if removeBroken && len(brokenEndpoints) > 0 {
		if err := applyRemoveBroken(brokenEndpoints); err != nil {
			return err
		}
	}

	if failed > 0 {
		return fmt.Errorf("%d forwarder(s) failed", failed)
	}
	return nil
}

func applyRemoveBroken(broken map[string][]string) error {
	cfgPath := config.ResolveFile(cfgFile)
	removed, err := config.RemoveBrokenEndpoints(cfgPath, broken)
	if err != nil {
		return fmt.Errorf("remove broken: %w", err)
	}
	if len(removed) == 0 {
		return nil
	}
	fmt.Println("\nRemoved from config:")
	for _, r := range removed {
		fmt.Printf("  - %s\n", r)
	}
	fmt.Printf("Config updated: %s\n", cfgPath)
	return nil
}

// embeddingStats returns the dimension count and L2 norm of the first embedding.
func embeddingStats(embeddings [][]float32) (dims int, norm float32) {
	if len(embeddings) == 0 || len(embeddings[0]) == 0 {
		return 0, 0
	}
	vec := embeddings[0]
	dims = len(vec)
	var sum float64
	for _, v := range vec {
		sum += float64(v) * float64(v)
	}
	return dims, float32(math.Sqrt(sum))
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n-3] + "..."
}
