package handler

import (
	"bytes"
	"encoding/json"
	"io"
	"net/http"
	"os"
	"strings"
)

var deepseekToken = os.Getenv("DEEPSEEK_TOKEN")

var deepseekModels = map[string]bool{
	"deepseek":                   true,
	"deepseek-expert":             true,
	"deepseek-r1":                 true,
	"deepseek-search":             true,
	"deepseek-expert-r1":          true,
	"deepseek-expert-search":      true,
	"deepseek-r1-search":          true,
	"deepseek-expert-r1-search":   true,
	"deepseek-v4":                 true,
	"deepseek-v4-flash":           true,
	"deepseek-v4-pro":             true,
}

func IsDeepSeekModel(model string) bool {
	_, ok := deepseekModels[strings.ToLower(model)]
	return ok
}

func HandleDeepSeek(w http.ResponseWriter, r *http.Request, body []byte, apiKey string) {
	var req map[string]interface{}
	json.Unmarshal(body, &req)
	newBody, _ := json.Marshal(req)
	proxyReq, _ := http.NewRequest("POST", "https://chat.deepseek.com/api/v1/chat/completions", bytes.NewReader(newBody))
	proxyReq.Header.Set("Authorization", "Bearer "+deepseekToken)
	proxyReq.Header.Set("Content-Type", "application/json")
	proxyReq.Header.Set("User-Agent", "Mozilla/5.0")
	proxyReq.Header.Set("Referer", "https://chat.deepseek.com/")
	client := &http.Client{}
	resp, err := client.Do(proxyReq)
	if err != nil { http.Error(w, "DeepSeek error", 500); return }
	defer resp.Body.Close()
	respBody, _ := io.ReadAll(resp.Body)
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(resp.StatusCode)
	w.Write(respBody)
	var result map[string]interface{}
	if json.Unmarshal(respBody, &result) == nil {
		if usage, ok := result["usage"].(map[string]interface{}); ok {
			if total, ok := usage["total_tokens"].(float64); ok {
				TrackUsage(apiKey, int(total))
			}
		}
	}
}