package handler

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

var openRouterKey = os.Getenv("OP_Rowproxy_key")

var openRouterModels = map[string]string{
	"deepseek-r1":        "deepseek/deepseek-r1:free",
	"deepseek-v3":        "deepseek/deepseek-v3.2:free",
	"deepseek-r1-0528":   "deepseek/deepseek-r1-0528:free",
	"deepseek-v4-flash":  "deepseek/deepseek-v4-flash:free",
	"deepseek-v4-pro":    "deepseek/deepseek-v4-pro:free",
}

func IsOpenRouterModel(model string) bool {
	_, ok := openRouterModels[model]
	return ok
}

func HandleOpenRouter(w http.ResponseWriter, r *http.Request, body []byte, apiKey string) {
	var req map[string]interface{}
	json.Unmarshal(body, &req)
	if m, ok := req["model"].(string); ok {
		if mapped, found := openRouterModels[m]; found {
			req["model"] = mapped
		}
	}
	newBody, _ := json.Marshal(req)
	proxyReq, _ := http.NewRequest("POST", "https://openrouter.ai/api/v1/chat/completions", bytes.NewReader(newBody))
	proxyReq.Header.Set("Authorization", "Bearer "+openRouterKey)
	proxyReq.Header.Set("Content-Type", "application/json")
	client := &http.Client{}
	resp, err := client.Do(proxyReq)
	if err != nil { http.Error(w, "OpenRouter 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))
			}
		}
	}
}