package main

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

	"github.com/google/uuid"
)

const (
	waveAPIURL = "https://cfapi.waveterm.dev/api/waveai"
)

type ResponsesRequest struct {
	Model           string                 `json:"model"`
	Input           json.RawMessage        `json:"input"`
	Instructions    string                 `json:"instructions,omitempty"`
	MaxOutputTokens int                    `json:"max_output_tokens,omitempty"`
	Temperature     float64                `json:"temperature,omitempty"`
	TopP            float64                `json:"top_p,omitempty"`
	Stream          *bool                  `json:"stream,omitempty"`
	Tools           json.RawMessage        `json:"tools,omitempty"`
	ToolChoice      interface{}            `json:"tool_choice,omitempty"`
	Reasoning       *ReasoningConfig       `json:"reasoning,omitempty"`
	Metadata        map[string]interface{} `json:"metadata,omitempty"`
}

type ReasoningConfig struct {
	Effort  string `json:"effort,omitempty"`
	Summary string `json:"summary,omitempty"`
}

type WaveRequest struct {
	Input           json.RawMessage  `json:"input"`
	Instructions    string           `json:"instructions,omitempty"`
	MaxOutputTokens int              `json:"max_output_tokens,omitempty"`
	Model           string           `json:"model"`
	Reasoning       *ReasoningConfig `json:"reasoning,omitempty"`
	Stream          bool             `json:"stream"`
	StreamOptions   *StreamOptions   `json:"stream_options,omitempty"`
	Text            *TextConfig      `json:"text,omitempty"`
	Tools           json.RawMessage  `json:"tools,omitempty"`
	ToolChoice      interface{}      `json:"tool_choice,omitempty"`
	Temperature     float64          `json:"temperature,omitempty"`
	TopP            float64          `json:"top_p,omitempty"`
}

type StreamOptions struct {
	IncludeObfuscation bool `json:"include_obfuscation"`
}

type TextConfig struct {
	Verbosity string `json:"verbosity,omitempty"`
}

func main() {
	port := os.Getenv("PORT")
	if port == "" {
		port = "7860"
	}

	http.HandleFunc("/v1/responses", handleResponses)
	http.HandleFunc("/responses", handleResponses)
	http.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(`{"status":"ok"}`))
	})

	log.Printf("Server starting on :%s", port)
	if err := http.ListenAndServe(":"+port, nil); err != nil {
		log.Fatalf("Server failed: %v", err)
	}
}

func handleResponses(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	body, _ := io.ReadAll(r.Body)
	var req ResponsesRequest
	json.Unmarshal(body, &req)

	waveReq := convertToWaveRequest(&req)
	stream := true
	if req.Stream != nil {
		stream = *req.Stream
	}

	if stream {
		handleStreamingResponse(w, waveReq)
	} else {
		handleNonStreamingResponse(w, waveReq)
	}
}

func convertToWaveRequest(req *ResponsesRequest) *WaveRequest {
	maxTokens := req.MaxOutputTokens
	if maxTokens == 0 {
		maxTokens = 4096
	}

	return &WaveRequest{
		Input:           req.Input,
		Instructions:    req.Instructions,
		MaxOutputTokens: maxTokens,
		Model:           "gpt-5.1",
		Stream:          true,
		StreamOptions:   &StreamOptions{IncludeObfuscation: false},
		Text:            &TextConfig{Verbosity: "low"},
		Reasoning:       &ReasoningConfig{Effort: "medium", Summary: "auto"},
	}
}

func makeWaveRequest(waveReq *WaveRequest) (*http.Response, error) {
	body, _ := json.Marshal(waveReq)
	httpReq, _ := http.NewRequest(http.MethodPost, waveAPIURL, bytes.NewReader(body))

	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Accept", "text/event-stream")
	httpReq.Header.Set("X-Wave-Apitype", "openai-responses")
	httpReq.Header.Set("X-Wave-Clientid", uuid.New().String())

	client := &http.Client{}
	return client.Do(httpReq)
}

func handleStreamingResponse(w http.ResponseWriter, waveReq *WaveRequest) {
	resp, err := makeWaveRequest(waveReq)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()

	w.Header().Set("Content-Type", "text/event-stream")
	flusher, _ := w.(http.Flusher)
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		w.Write([]byte(scanner.Text() + "\n"))
		flusher.Flush()
	}
}

func handleNonStreamingResponse(w http.ResponseWriter, waveReq *WaveRequest) {
	resp, _ := makeWaveRequest(waveReq)
	defer resp.Body.Close()

	var outputText strings.Builder
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.HasPrefix(line, "data: ") {
			var event map[string]interface{}
			json.Unmarshal([]byte(strings.TrimPrefix(line, "data: ")), &event)
			if delta, ok := event["delta"].(string); ok {
				outputText.WriteString(delta)
			}
		}
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{
		"choices": []map[string]interface{}{
			{"message": map[string]string{"content": outputText.String()}},
		},
	})
}