package handler

import (
	"encoding/json"
	"fmt"
	"net/http"

	"github.com/edgeai/gateway/internal/adapter"
	"github.com/edgeai/gateway/pkg/api"
)

// SSEWriter writes Server-Sent Events to an HTTP response.
type SSEWriter struct {
	w       http.ResponseWriter
	flusher http.Flusher
}

// NewSSEWriter creates a new SSEWriter. Returns nil if streaming is not supported.
func NewSSEWriter(w http.ResponseWriter) *SSEWriter {
	flusher, ok := w.(http.Flusher)
	if !ok {
		return nil
	}

	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("X-Accel-Buffering", "no")

	return &SSEWriter{w: w, flusher: flusher}
}

// WriteChunk writes a single SSE data event.
func (s *SSEWriter) WriteChunk(data any) error {
	jsonData, err := json.Marshal(data)
	if err != nil {
		return fmt.Errorf("marshal sse data: %w", err)
	}
	fmt.Fprintf(s.w, "data: %s\n\n", jsonData)
	s.flusher.Flush()
	return nil
}

// WriteDone writes the [DONE] marker.
func (s *SSEWriter) WriteDone() {
	fmt.Fprintf(s.w, "data: [DONE]\n\n")
	s.flusher.Flush()
}

// StreamChatCompletion streams chunks from an adapter to the client in OpenAI SSE format.
func StreamChatCompletion(sse *SSEWriter, ch <-chan adapter.StreamChunk, requestID, taskID, model string) (int, int, error) {
	inputTokens := 0
	outputTokens := 0

	for chunk := range ch {
		if chunk.Error != nil {
			return inputTokens, outputTokens, chunk.Error
		}

		if chunk.Done {
			if chunk.InputTokens > 0 {
				inputTokens = chunk.InputTokens
			}
			if chunk.OutputTokens > 0 {
				outputTokens = chunk.OutputTokens
			}

			// Write final chunk with finish_reason
			sseChunk := map[string]any{
				"id":      requestID,
				"object":  "chat.completion.chunk",
				"model":   model,
				"choices": []map[string]any{
					{
						"index":         0,
						"delta":         map[string]any{},
						"finish_reason": chunk.FinishReason,
					},
				},
			}
			if inputTokens > 0 || outputTokens > 0 {
				sseChunk["usage"] = map[string]int{
					"input_tokens":  inputTokens,
					"output_tokens": outputTokens,
					"total_tokens":  inputTokens + outputTokens,
				}
			}
			sse.WriteChunk(sseChunk)
			sse.WriteDone()
			return inputTokens, outputTokens, nil
		}

		// Write content delta
		sseChunk := map[string]any{
			"id":      requestID,
			"object":  "chat.completion.chunk",
			"model":   model,
			"choices": []map[string]any{
				{
					"index": 0,
					"delta": map[string]any{
						"content": chunk.Delta,
					},
					"finish_reason": nil,
				},
			},
		}
		sse.WriteChunk(sseChunk)
	}

	return inputTokens, outputTokens, nil
}

// BuildChatResponse creates a non-streaming ChatResponse from adapter result.
func BuildChatResponse(requestID, taskID, logicalModel string, resp *adapter.ChatResponse) api.ChatResponse {
	return api.ChatResponse{
		RequestID:   requestID,
		TaskID:      taskID,
		Status:      "completed",
		Model:       logicalModel,
		Choices: []api.Choice{
			{
				Index: 0,
				Message: &api.Message{
					Role:    "assistant",
					Content: resp.Content,
				},
				FinishReason: resp.FinishReason,
			},
		},
		LogicalModel: logicalModel,
		ActualModel:  resp.ActualModel,
		Usage: &api.Usage{
			InputTokens:  resp.InputTokens,
			OutputTokens: resp.OutputTokens,
			TotalTokens:  resp.InputTokens + resp.OutputTokens,
		},
	}
}