package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/sugarme/tokenizer"
	"github.com/sugarme/tokenizer/pretrained"
	"github.com/yalue/onnxruntime_go"
)

var (
	modelPath     = "/app/model.onnx"
	apiSecret     = os.Getenv("API_SECRET")
	tok           *tokenizer.Tokenizer
	session       *onnxruntime_go.DynamicAdvancedSession
	maxLen        = int64(128)
	embeddingSize = int64(384)
	version       = "dev"
)

// bge-small-en-v1.5 uses asymmetric query/passage embeddings: queries are
// prefixed with this instruction, passages (package descriptions) are not.
const queryInstruction = "Represent this sentence for searching relevant passages: "

func init() {
	// Initialize ONNX Runtime
	libPath := os.Getenv("ONNXRUNTIME_SHARED_LIBRARY_PATH")
	if libPath == "" {
		libPath = "/usr/local/lib/libonnxruntime.so"
	}
	onnxruntime_go.SetSharedLibraryPath(libPath)

	err := onnxruntime_go.InitializeEnvironment()
	if err != nil {
		log.Fatalf("Failed to initialize ONNX Runtime: %v", err)
	}

	// Load tokenizer
	tok, err = pretrained.FromFile("/app/tokenizer.json")
	if err != nil {
		log.Fatalf("Failed to load tokenizer: %v", err)
	}

	// Load ONNX model. bge-small's export already provides a pooled,
	// L2-normalized "sentence_embedding" output, so no manual pooling.
	session, err = onnxruntime_go.NewDynamicAdvancedSession(
		modelPath,
		[]string{"input_ids", "attention_mask", "token_type_ids"},
		[]string{"sentence_embedding"},
		nil)
	if err != nil {
		log.Fatalf("Failed to load ONNX model: %v", err)
	}

	log.Println("Model and tokenizer loaded. Starting server...")
}

func encode(text string) ([]int64, []int64, []int64) {
	// Tokenize using the proper tokenizer
	inputSeq := tokenizer.NewInputSequence(text)
	input := tokenizer.NewSingleEncodeInput(inputSeq)
	encoding, err := tok.Encode(input, true)
	if err != nil {
		log.Fatalf("Failed to tokenize: %v", err)
	}

	inputIds := make([]int64, len(encoding.GetIds()))
	for i, id := range encoding.GetIds() {
		inputIds[i] = int64(id)
	}

	// Truncate or pad
	paddedIds := make([]int64, maxLen)
	copy(paddedIds, inputIds)

	attentionMask := make([]int64, maxLen)
	tokenTypeIds := make([]int64, maxLen)
	for i := 0; i < int(maxLen); i++ {
		if i < len(inputIds) {
			attentionMask[i] = 1
		} else {
			attentionMask[i] = 0
		}
		tokenTypeIds[i] = 0
	}

	return paddedIds, attentionMask, tokenTypeIds
}

func vectorHandler(w http.ResponseWriter, r *http.Request) {
	// Check API secret
	if apiSecret != "" {
		auth := r.Header.Get("Authorization")
		if auth != "Bearer "+apiSecret {
			http.Error(w, "Unauthorized - Invalid or missing API key", http.StatusUnauthorized)
			return
		}
	}

	if r.Method != "POST" {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	var req struct {
		Text string `json:"text"`
		Type string `json:"type"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Bad request", http.StatusBadRequest)
		return
	}

	if req.Text == "" {
		http.Error(w, "Text is required", http.StatusBadRequest)
		return
	}

	// bge expects queries to be prefixed with the retrieval instruction;
	// passages (package descriptions) are embedded as-is. The default is
	// "query" so existing callers sending a bare {"text": ...} keep
	// working unchanged; documents must opt out with type=passage.
	text := req.Text
	if req.Type != "passage" {
		text = queryInstruction + text
	}

	// Encode text
	inputIds, attentionMask, tokenTypeIds := encode(text)

	// Create input tensors
	inputShape := onnxruntime_go.NewShape(1, maxLen)

	inputIdsTensor, err := onnxruntime_go.NewTensor(inputShape, inputIds)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to create input_ids tensor: %v", err), http.StatusInternalServerError)
		return
	}
	defer inputIdsTensor.Destroy()

	attentionMaskTensor, err := onnxruntime_go.NewTensor(inputShape, attentionMask)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to create attention_mask tensor: %v", err), http.StatusInternalServerError)
		return
	}
	defer attentionMaskTensor.Destroy()

	tokenTypeIdsTensor, err := onnxruntime_go.NewTensor(inputShape, tokenTypeIds)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to create token_type_ids tensor: %v", err), http.StatusInternalServerError)
		return
	}
	defer tokenTypeIdsTensor.Destroy()

	// Create output tensor: "sentence_embedding" has shape [1, 384].
	outputShape := onnxruntime_go.NewShape(1, embeddingSize)
	outputTensor, err := onnxruntime_go.NewEmptyTensor[float32](outputShape)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failed to create output tensor: %v", err), http.StatusInternalServerError)
		return
	}
	defer outputTensor.Destroy()

	// Run inference
	inputs := []onnxruntime_go.Value{inputIdsTensor, attentionMaskTensor, tokenTypeIdsTensor}
	outputs := []onnxruntime_go.Value{outputTensor}

	err = session.Run(inputs, outputs)
	if err != nil {
		http.Error(w, fmt.Sprintf("Inference error: %v", err), http.StatusInternalServerError)
		return
	}

	// The sentence embedding is already mean-pooled and L2-normalized.
	embeddings := outputTensor.GetData()
	vector := make([]float32, embeddingSize)
	for i := 0; i < int(embeddingSize); i++ {
		vector[i] = embeddings[i]
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{
		"vector": vector,
	})
}

type statusRecorder struct {
	http.ResponseWriter
	status int
}

func (r *statusRecorder) WriteHeader(code int) {
	r.status = code
	r.ResponseWriter.WriteHeader(code)
}

func loggingMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		rec := &statusRecorder{ResponseWriter: w, status: http.StatusOK}
		next.ServeHTTP(rec, r)
		log.Printf("%s %s %s %d %s", r.RemoteAddr, r.Method, r.URL.Path, rec.status, time.Since(start))
	})
}

func healthHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
}

func versionHandler(w http.ResponseWriter, r *http.Request) {
	v := version
	if v == "" || v == "dev" {
		if e := os.Getenv("SOURCE_COMMIT"); e != "" {
			v = e
		}
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]string{"version": v})
}

func main() {
	port := os.Getenv("PORT")
	if port == "" {
		port = "8080"
	}

	http.HandleFunc("/vector", vectorHandler)
	http.HandleFunc("/ok", healthHandler)
	http.HandleFunc("/version", versionHandler)
	log.SetOutput(os.Stdout)
	log.Printf("Starting server on port %s...\n", port)
	log.Fatal(http.ListenAndServe(":"+port, loggingMiddleware(http.DefaultServeMux)))
}