// mockupstream 本地 mock OpenAI 上游服务（联调代理链路，无需真实 key）。
// 支持 /v1/chat/completions 与 /v1/responses，含流式与非流式。
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"time"
)

func main() {
	addr := flag.String("addr", ":9000", "listen address")
	flag.Parse()

	http.HandleFunc("/v1/chat/completions", func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		var req struct {
			Model  string `json:"model"`
			Stream bool   `json:"stream"`
		}
		_ = json.Unmarshal(body, &req)
		if req.Stream {
			streamChat(w, req.Model)
			return
		}
		replyChat(w, req.Model)
	})

	http.HandleFunc("/v1/responses", func(w http.ResponseWriter, r *http.Request) {
		body, _ := io.ReadAll(r.Body)
		var req struct {
			Model  string `json:"model"`
			Stream bool   `json:"stream"`
		}
		_ = json.Unmarshal(body, &req)
		if req.Stream {
			streamResponses(w, req.Model)
			return
		}
		replyResponses(w, req.Model)
	})

	http.HandleFunc("/v1/models", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, `{"object":"list","data":[{"id":"gpt-4o-mini","object":"model"},{"id":"gpt-4o","object":"model"}]}`)
	})

	log.Printf("mock upstream listening on %s", *addr)
	log.Fatal(http.ListenAndServe(*addr, nil))
}

func replyChat(w http.ResponseWriter, model string) {
	w.Header().Set("Content-Type", "application/json")
	resp := map[string]any{
		"id":      "chatcmpl-mock123",
		"object":  "chat.completion",
		"model":   model,
		"choices": []any{map[string]any{
			"index": 0,
			"message": map[string]any{"role": "assistant", "content": "你好，这是 mock 上游的回复。"},
			"finish_reason": "stop",
		}},
		"usage": map[string]any{"prompt_tokens": 12, "completion_tokens": 9, "total_tokens": 21},
	}
	_ = json.NewEncoder(w).Encode(resp)
}

func streamChat(w http.ResponseWriter, model string) {
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	fl, _ := w.(http.Flusher)
	chunks := []string{"你好", "，这是", " mock ", "流式回复。"}
	for i, c := range chunks {
		chunk := map[string]any{
			"id": "chatcmpl-mock123", "object": "chat.completion.chunk", "model": model,
			"choices": []any{map[string]any{
				"index": 0,
				"delta": map[string]any{"content": c},
				"finish_reason": nil,
			}},
		}
		if i == len(chunks)-1 {
			chunk["choices"] = []any{map[string]any{
				"index": 0, "delta": map[string]any{}, "finish_reason": "stop",
			}}
		}
		b, _ := json.Marshal(chunk)
		fmt.Fprintf(w, "data: %s\n\n", b)
		fl.Flush()
		time.Sleep(50 * time.Millisecond)
	}
	// usage 末块（include_usage 时返回）
	usage := map[string]any{
		"id": "chatcmpl-mock123", "object": "chat.completion.chunk", "model": model,
		"choices": []any{}, "usage": map[string]any{"prompt_tokens": 12, "completion_tokens": 9, "total_tokens": 21},
	}
	b, _ := json.Marshal(usage)
	fmt.Fprintf(w, "data: %s\n\n", b)
	fmt.Fprint(w, "data: [DONE]\n\n")
	fl.Flush()
}

func replyResponses(w http.ResponseWriter, model string) {
	w.Header().Set("Content-Type", "application/json")
	resp := map[string]any{
		"id":      "resp_mock456",
		"object":  "response",
		"model":   model,
		"status":  "completed",
		"output":  []any{map[string]any{
			"type": "message", "role": "assistant",
			"content": []any{map[string]any{"type": "output_text", "text": "这是 Responses API 的 mock 回复。"}},
		}},
		"usage": map[string]any{"input_tokens": 15, "output_tokens": 11, "total_tokens": 26},
	}
	_ = json.NewEncoder(w).Encode(resp)
}

func streamResponses(w http.ResponseWriter, model string) {
	w.Header().Set("Content-Type", "text/event-stream")
	fl, _ := w.(http.Flusher)
	events := []map[string]any{
		{"type": "response.created", "response": map[string]any{"id": "resp_mock456", "model": model, "object": "response", "status": "in_progress"}},
		{"type": "response.output_text.delta", "delta": "流式 Responses 内容", "item_id": "msg_1", "output_index": 0, "content_index": 0},
		{"type": "response.completed", "response": map[string]any{
			"id": "resp_mock456", "object": "response", "model": model, "status": "completed",
			"usage": map[string]any{"input_tokens": 15, "output_tokens": 11, "total_tokens": 26},
		}},
	}
	for _, e := range events {
		b, _ := json.Marshal(e)
		fmt.Fprintf(w, "event: %s\ndata: %s\n\n", e["type"], b)
		fl.Flush()
		time.Sleep(50 * time.Millisecond)
	}
}