Local_Perplexity/internal/handler/proxy.go

package handler

import (
	"encoding/json"
	"log"
	"net/http"
	"strings"

	"ai-platform/internal/model"
	"ai-platform/internal/service"
)

// ProxyHandler — HTTP-хендлеры для проксирования запросов к Ollama
type ProxyHandler struct {
	ollamaClient *service.OllamaClient
	router       *service.Router
}

func NewProxyHandler(ollamaClient *service.OllamaClient, router *service.Router) *ProxyHandler {
	return &ProxyHandler{
		ollamaClient: ollamaClient,
		router:       router,
	}
}

// HandleChat — POST /api/chat
// Декодирует запрос → маршрутизация → подмена модели → стриминг ответа
func (h *ProxyHandler) HandleChat(w http.ResponseWriter, r *http.Request) {
	var req model.ChatRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, `{"error":"invalid request body"}`, http.StatusBadRequest)
		return
	}

	// Маршрутизация: "auto"/пустая → Router LLM, иначе → как есть
	targetModel := h.router.Route(r.Context(),req.Model, req.Messages)
	req.Model = targetModel

	log.Printf("proxy: /api/chat → model=%s", targetModel)

	if err := h.ollamaClient.ProxyChat(r.Context(), w, req); err != nil {
		log.Printf("proxy: chat error: %v", err)
		// Если ещё не начали писать ответ, вернём ошибку
	}
}

// HandleGenerate — POST /api/generate
// Аналогично HandleChat, но для generate-эндпоинта
func (h *ProxyHandler) HandleGenerate(w http.ResponseWriter, r *http.Request) {
	var req model.GenerateRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, `{"error":"invalid request body"}`, http.StatusBadRequest)
		return
	}

	// Для generate маршрутизация по модели (без анализа сообщений)
	if req.Model == "" || req.Model == "auto" {
		msgs := []model.Message{{Role: "user", Content: req.Prompt}}
		req.Model = h.router.Route(r.Context(),req.Model, msgs)
	}

	log.Printf("proxy: /api/generate → model=%s", req.Model)

	if err := h.ollamaClient.ProxyGenerate(r.Context(), w, req); err != nil {
		log.Printf("proxy: generate error: %v", err)
	}
}

// HandleTags — GET /api/tags
// Проксирует список моделей из Ollama
func (h *ProxyHandler) HandleTags(w http.ResponseWriter, r *http.Request) {
	tags, err := h.ollamaClient.GetTags(r.Context())
	if err != nil {
		log.Printf("proxy: tags error: %v", err)
		http.Error(w, `{"error":"failed to get models"}`, http.StatusBadGateway)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(tags)
}

// HandleChatV1 — POST /v1/chat/completions (OpenAI-совместимый формат)
// Используется Codex CLI и другими OpenAI-совместимыми клиентами
func (h *ProxyHandler) HandleChatV1(w http.ResponseWriter, r *http.Request) {
	var req model.ChatRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, `{"error":{"message":"invalid request body"}}`, http.StatusBadRequest)
		return
	}

	targetModel := h.router.Route(r.Context(),req.Model, req.Messages)
	req.Model = targetModel

	log.Printf("proxy: /v1/chat/completions → model=%s", targetModel)

	if err := h.ollamaClient.ProxyChatV1(r.Context(), w, req); err != nil {
		log.Printf("proxy: v1 chat error: %v", err)
	}
}

// HandleResponses — POST /responses и /v1/responses (Codex CLI)
func (h *ProxyHandler) HandleResponses(w http.ResponseWriter, r *http.Request) {
	var req model.ResponsesRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, `{"error":{"message":"invalid request body"}}`, http.StatusBadRequest)
		return
	}

	allMessages := req.ToMessages()

	// Маршрутизация по всем сообщениям (включая системные, для определения контекста)
	targetModel := h.router.Route(r.Context(),req.Model, allMessages)

	// Фильтруем сообщения для Ollama: убираем системный мусор Codex CLI,
	// конвертируем developer→system, оставляем только user/assistant/system
	messages := filterMessagesForOllama(allMessages)

	log.Printf("proxy: /responses → model=%s (total=%d, filtered=%d)", targetModel, len(allMessages), len(messages))

	if err := h.ollamaClient.ProxyResponsesV1(r.Context(), w, messages, targetModel); err != nil {
		log.Printf("proxy: responses error: %v", err)
	}
}

// filterMessagesForOllama убирает системный мусор Codex CLI и оставляет только
// релевантные сообщения для Ollama (user, assistant, один system).
func filterMessagesForOllama(messages []model.Message) []model.Message {
	var filtered []model.Message
	hasSystem := false

	for _, m := range messages {
		// Пропускаем developer-сообщения (Codex CLI permissions, collaboration mode)
		if m.Role == "developer" {
			continue
		}

		// Пропускаем системные инструкции Codex CLI (длинные system-сообщения)
		if m.Role == "system" && len(m.Content) > 500 {
			continue
		}

		// Пропускаем user-сообщения которые являются системными инструкциями
		if m.Role == "user" && isCodexBoilerplate(m.Content) {
			continue
		}

		// Оставляем только один system-сообщение
		if m.Role == "system" {
			if hasSystem {
				continue
			}
			hasSystem = true
		}

		filtered = append(filtered, m)
	}

	// Если после фильтрации нет сообщений — вернуть оригинал
	if len(filtered) == 0 {
		return messages
	}

	return filtered
}

// isCodexBoilerplate проверяет, является ли user-сообщение системными инструкциями Codex CLI
func isCodexBoilerplate(content string) bool {
	if len(content) > 500 {
		lower := strings.ToLower(content[:500])
		return strings.Contains(lower, "<instructions>") ||
			strings.Contains(lower, "<permissions") ||
			strings.Contains(lower, "<environment_context>") ||
			strings.Contains(lower, "agents.md")
	}
	return false
}

// HandleGetResponse — GET /responses и /responses/{id}
// Codex CLI может запрашивать статус ответа
func (h *ProxyHandler) HandleGetResponse(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	w.Write([]byte(`{"id":"resp_stub","object":"response","status":"completed"}`))
}

// HandleModelsV1 — GET /v1/models (OpenAI-совместимый формат)
func (h *ProxyHandler) HandleModelsV1(w http.ResponseWriter, r *http.Request) {
	body, err := h.ollamaClient.GetModelsV1(r.Context())
	if err != nil {
		log.Printf("proxy: v1 models error: %v", err)
		http.Error(w, `{"error":{"message":"failed to get models"}}`, http.StatusBadGateway)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	w.Write(body)
}