package openai

import (
	"net/http"
	"strings"
	"time"

	edgeservice "iop/apps/edge/internal/service"
	iop "iop/proto/gen/iop"
)

var anthropicResponseHeaderAllowlist = map[string]struct{}{
	"Cache-Control":  {},
	"Content-Length": {},
	"Content-Type":   {},
	"Request-Id":     {},
	"Retry-After":    {},
	"X-Request-Id":   {},
	"X-Robots-Tag":   {},
}

func (s *Server) writeAnthropicNativeTunnelResponse(w http.ResponseWriter, r *http.Request, handle edgeservice.ProviderTunnelResult) {
	frames := handle.Stream().Frames
	if frames == nil {
		writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider tunnel is unavailable")
		return
	}
	flusher, _ := w.(http.Flusher)
	timer := time.NewTimer(handle.WaitTimeout())
	defer timer.Stop()
	wroteHeader := false

	for {
		select {
		case <-r.Context().Done():
			s.cancelRunOnHTTPGiveUp(handle.Dispatch(), r.Context().Err())
			return
		case <-timer.C:
			s.cancelRunOnHTTPGiveUp(handle.Dispatch(), errRunTimedOut)
			if !wroteHeader {
				writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider response timed out")
			}
			return
		case frame, ok := <-frames:
			if !ok {
				if !wroteHeader {
					writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider tunnel closed before a response")
				}
				return
			}
			switch frame.GetKind() {
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START:
				if wroteHeader {
					continue
				}
				copyAnthropicResponseHeaders(w.Header(), frame.GetHeaders())
				status := int(frame.GetStatusCode())
				if status == 0 {
					status = http.StatusOK
				}
				w.WriteHeader(status)
				wroteHeader = true
				if flusher != nil {
					flusher.Flush()
				}
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY:
				if len(frame.GetBody()) == 0 {
					continue
				}
				if !wroteHeader {
					w.Header().Set("Content-Type", "application/json")
					w.WriteHeader(http.StatusOK)
					wroteHeader = true
				}
				if _, err := w.Write(frame.GetBody()); err != nil {
					s.sendCancelRun(handle.Dispatch())
					return
				}
				if flusher != nil {
					flusher.Flush()
				}
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR:
				if !wroteHeader {
					writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider tunnel failed")
				} else if strings.Contains(strings.ToLower(w.Header().Get("Content-Type")), "text/event-stream") {
					_ = writeAnthropicSSEEvent(w, "error", anthropicErrorResponse{
						Type: "error", Error: errorBody{Type: "api_error", Message: "provider tunnel failed"},
					})
					if flusher != nil {
						flusher.Flush()
					}
				}
				return
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END:
				if !wroteHeader {
					writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider tunnel ended before a response")
				}
				return
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE:
				continue
			}
		}
	}
}

func copyAnthropicResponseHeaders(dst http.Header, headers map[string]string) {
	for key, value := range headers {
		canonical := http.CanonicalHeaderKey(key)
		_, exact := anthropicResponseHeaderAllowlist[canonical]
		lower := strings.ToLower(key)
		if !exact && !strings.HasPrefix(lower, "anthropic-ratelimit-") && !strings.HasPrefix(lower, "ratelimit-") {
			continue
		}
		dst.Set(canonical, value)
	}
}