main.go

package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"sort"
	"strings"
	"time"
)

const (
	NvidiaBaseURL = "https://integrate.api.nvidia.com/v1"
	NvidiaAPIKey  = "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"
	GatewayAPIKey = "connect"
)

var modelAliases = map[string]string{
	"Bielik-11b":      "speakleash/bielik-11b-v2.6-instruct",
	"GLM-4.7":         "z-ai/glm4.7",
	"Mistral-Small-4": "mistralai/mistral-small-4-119b-2603",
	"DeepSeek-V3.1":   "deepseek-ai/deepseek-v3.1",
	"Kimi-K2":         "moonshotai/kimi-k2-instruct",
}

type Message struct {
	Role       string      `json:"role"`
	Content    interface{} `json:"content"`
	ToolCallID string      `json:"tool_call_id,omitempty"`
	ToolCalls  interface{} `json:"tool_calls,omitempty"`
	Name       string      `json:"name,omitempty"`
}

type ChatRequest struct {
	Model       string        `json:"model"`
	Messages    []Message     `json:"messages"`
	Stream      *bool         `json:"stream,omitempty"`
	Tools       []interface{} `json:"tools,omitempty"`
	ToolChoice  interface{}   `json:"tool_choice,omitempty"`
	Temperature *float64      `json:"temperature,omitempty"`
	MaxTokens   *int          `json:"max_tokens,omitempty"`
	TopP        *float64      `json:"top_p,omitempty"`
	Stop        interface{}   `json:"stop,omitempty"`
}

type UpstreamRequest struct {
	Model       string                 `json:"model"`
	Messages    []Message              `json:"messages"`
	Stream      bool                   `json:"stream"`
	Tools       []interface{}          `json:"tools,omitempty"`
	ToolChoice  interface{}            `json:"tool_choice,omitempty"`
	Temperature *float64               `json:"temperature,omitempty"`
	MaxTokens   *int                   `json:"max_tokens,omitempty"`
	TopP        *float64               `json:"top_p,omitempty"`
	Stop        interface{}            `json:"stop,omitempty"`
	ExtraBody   map[string]interface{} `json:"extra_body,omitempty"`
}

type RawChunk struct {
	ID      string      `json:"id"`
	Object  string      `json:"object"`
	Created int64       `json:"created"`
	Model   string      `json:"model"`
	Choices []RawChoice `json:"choices"`
	Usage   interface{} `json:"usage,omitempty"`
}

type RawChoice struct {
	Index        int      `json:"index"`
	Delta        RawDelta `json:"delta"`
	FinishReason *string  `json:"finish_reason"`
}

type RawDelta struct {
	Role      string       `json:"role,omitempty"`
	Content   *string      `json:"content,omitempty"`
	ToolCalls []RawToolCall `json:"tool_calls,omitempty"`
}

type RawToolCall struct {
	Index    int         `json:"index"`
	ID       string      `json:"id,omitempty"`
	Type     string      `json:"type,omitempty"`
	Function RawFunction `json:"function"`
}

type RawFunction struct {
	Name      string `json:"name,omitempty"`
	Arguments string `json:"arguments,omitempty"`
}

type AccumToolCall struct {
	Index int
	ID    string
	Type  string
	Name  string
	Args  string
}

func resolveModel(requested string) string {
	if full, ok := modelAliases[requested]; ok {
		return full
	}
	for _, full := range modelAliases {
		if full == requested {
			return requested
		}
	}
	return requested
}

func injectSystemPrompt(messages []Message, modelID string) []Message {
	filtered := make([]Message, 0, len(messages))
	for _, m := range messages {
		if m.Role != "system" {
			filtered = append(filtered, m)
		}
	}
	prompt, ok := systemPrompts[modelID]
	if !ok || prompt == "" {
		return filtered
	}
	return append([]Message{{Role: "system", Content: prompt}}, filtered...)
}

func authenticate(r *http.Request) bool {
	auth := r.Header.Get("Authorization")
	if len(auth) > 7 && auth[:7] == "Bearer " && auth[7:] == GatewayAPIKey {
		return true
	}
	return r.Header.Get("x-api-key") == GatewayAPIKey
}

func handleModels(w http.ResponseWriter, r *http.Request) {
	if !authenticate(r) {
		http.Error(w, `{"error":{"message":"Unauthorized"}}`, http.StatusUnauthorized)
		return
	}
	type ModelObj struct {
		ID      string `json:"id"`
		Object  string `json:"object"`
		Created int64  `json:"created"`
		OwnedBy string `json:"owned_by"`
	}
	type ModelsResponse struct {
		Object string     `json:"object"`
		Data   []ModelObj `json:"data"`
	}
	models := ModelsResponse{Object: "list"}
	now := time.Now().Unix()
	for alias := range modelAliases {
		models.Data = append(models.Data, ModelObj{ID: alias, Object: "model", Created: now, OwnedBy: "nvidia"})
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(models)
}

func handleBaseURL(w http.ResponseWriter, r *http.Request) {
	host := os.Getenv("SPACE_HOST")
	if host == "" {
		host = r.Host
	}
	w.Header().Set("Content-Type", "application/json")
	fmt.Fprintf(w, `{"url":"https://%s/v1"}`, host)
}

func handleChat(w http.ResponseWriter, r *http.Request) {
	if !authenticate(r) {
		http.Error(w, `{"error":{"message":"Unauthorized"}}`, http.StatusUnauthorized)
		return
	}
	if r.Method != http.MethodPost {
		http.Error(w, `{"error":{"message":"Method not allowed"}}`, http.StatusMethodNotAllowed)
		return
	}

	var req ChatRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, `{"error":{"message":"Invalid request body"}}`, http.StatusBadRequest)
		return
	}

	modelID := resolveModel(req.Model)
	upstream := UpstreamRequest{
		Model:       modelID,
		Messages:    injectSystemPrompt(req.Messages, modelID),
		Stream:      true,
		Tools:       req.Tools,
		ToolChoice:  req.ToolChoice,
		Temperature: req.Temperature,
		MaxTokens:   req.MaxTokens,
		TopP:        req.TopP,
		Stop:        req.Stop,
	}
	if modelID == "z-ai/glm4.7" {
		upstream.ExtraBody = map[string]interface{}{
			"chat_template_kwargs": map[string]interface{}{
				"enable_thinking": false,
			},
		}
	}

	body, err := json.Marshal(upstream)
	if err != nil {
		http.Error(w, `{"error":{"message":"Failed to marshal request"}}`, http.StatusInternalServerError)
		return
	}

	upstreamReq, err := http.NewRequest(http.MethodPost, NvidiaBaseURL+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		http.Error(w, `{"error":{"message":"Failed to create upstream request"}}`, http.StatusInternalServerError)
		return
	}
	upstreamReq.Header.Set("Content-Type", "application/json")
	upstreamReq.Header.Set("Authorization", "Bearer "+NvidiaAPIKey)
	upstreamReq.Header.Set("Accept", "text/event-stream")

	client := &http.Client{Timeout: 300 * time.Second}
	resp, err := client.Do(upstreamReq)
	if err != nil {
		http.Error(w, fmt.Sprintf(`{"error":{"message":"%s"}}`, err.Error()), http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		upstreamBody, _ := io.ReadAll(resp.Body)
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(resp.StatusCode)
		w.Write(upstreamBody)
		return
	}

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("X-Accel-Buffering", "no")
	w.WriteHeader(http.StatusOK)

	flusher, canFlush := w.(http.Flusher)
	emit := func(s string) {
		fmt.Fprint(w, s)
		if canFlush {
			flusher.Flush()
		}
	}

	scanner := bufio.NewScanner(resp.Body)
	scanner.Buffer(make([]byte, 1024*1024), 1024*1024)

	// Accumulate tool_calls across delta chunks, stream content chunks immediately
	accum := make(map[int]*AccumToolCall)
	var lastChunk RawChunk

	for scanner.Scan() {
		line := scanner.Text()

		if !strings.HasPrefix(line, "data: ") {
			emit(line + "\n")
			continue
		}

		data := strings.TrimPrefix(line, "data: ")

		if data == "[DONE]" {
			emit("data: [DONE]\n\n")
			continue
		}

		var chunk RawChunk
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			emit(line + "\n")
			continue
		}
		lastChunk = chunk

		isToolChunk := false
		isFinishToolCalls := false

		for _, choice := range chunk.Choices {
			if len(choice.Delta.ToolCalls) > 0 {
				isToolChunk = true
				for _, tc := range choice.Delta.ToolCalls {
					acc, ok := accum[tc.Index]
					if !ok {
						acc = &AccumToolCall{Index: tc.Index}
						accum[tc.Index] = acc
					}
					if tc.ID != "" {
						acc.ID = tc.ID
					}
					if tc.Type != "" {
						acc.Type = tc.Type
					}
					acc.Name += tc.Function.Name
					acc.Args += tc.Function.Arguments
				}
			}
			if choice.FinishReason != nil && *choice.FinishReason == "tool_calls" {
				isFinishToolCalls = true
			}
		}

		if isFinishToolCalls {
			// Emit one complete tool_calls chunk with all assembled tool calls
			indices := make([]int, 0, len(accum))
			for idx := range accum {
				indices = append(indices, idx)
			}
			sort.Ints(indices)

			assembled := make([]map[string]interface{}, 0, len(indices))
			for _, idx := range indices {
				acc := accum[idx]
				tcType := acc.Type
				if tcType == "" {
					tcType = "function"
				}
				assembled = append(assembled, map[string]interface{}{
					"index": idx,
					"id":    acc.ID,
					"type":  tcType,
					"function": map[string]string{
						"name":      acc.Name,
						"arguments": acc.Args,
					},
				})
			}

			fr := "tool_calls"
			out, _ := json.Marshal(map[string]interface{}{
				"id":      lastChunk.ID,
				"object":  "chat.completion.chunk",
				"created": lastChunk.Created,
				"model":   req.Model,
				"choices": []map[string]interface{}{
					{
						"index": 0,
						"delta": map[string]interface{}{
							"role":       "assistant",
							"content":    nil,
							"tool_calls": assembled,
						},
						"finish_reason": fr,
					},
				},
			})
			emit("data: " + string(out) + "\n\n")
			accum = make(map[int]*AccumToolCall)
			continue
		}

		// Skip intermediate tool_call delta chunks (already accumulating)
		if isToolChunk {
			continue
		}

		// Regular content chunk — stream immediately as-is
		emit("data: " + data + "\n\n")
	}
}

func loggingMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		log.Printf("[%s] %s %s", r.Method, r.URL.Path, r.RemoteAddr)
		next(w, r)
		log.Printf("[%s] %s done in %s", r.Method, r.URL.Path, time.Since(start))
	}
}

func corsMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization, x-api-key")
		if r.Method == http.MethodOptions {
			w.WriteHeader(http.StatusNoContent)
			return
		}
		next(w, r)
	}
}

func main() {
	port := os.Getenv("PORT")
	if port == "" {
		port = "7860"
	}
	mux := http.NewServeMux()
	mux.HandleFunc("/v1/chat/completions", corsMiddleware(loggingMiddleware(handleChat)))
	mux.HandleFunc("/v1/models", corsMiddleware(loggingMiddleware(handleModels)))
	mux.HandleFunc("/v1/base-url", corsMiddleware(handleBaseURL))
	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(`{"status":"ok"}`))
	})
	log.Printf("Gateway starting on :%s", port)
	if err := http.ListenAndServe(":"+port, mux); err != nil {
		log.Fatal(err)
	}
}