package openai

import (
	"bytes"
	"encoding/json"
	"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, publicModelID string) {
	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
	receivedResponseStart := false
	responseStatus := http.StatusOK
	responseStreaming := false
	rewriteResponse := strings.TrimSpace(publicModelID) != ""
	var responseBody []byte
	var streamRewriter *anthropicNativeModelRewriter

	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 receivedResponseStart || wroteHeader {
					continue
				}
				receivedResponseStart = true
				copyAnthropicResponseHeaders(w.Header(), frame.GetHeaders())
				status := int(frame.GetStatusCode())
				if status == 0 {
					status = http.StatusOK
				}
				responseStatus = status
				if rewriteResponse && status >= http.StatusOK && status < http.StatusMultipleChoices {
					w.Header().Del("Content-Length")
					responseStreaming = strings.Contains(strings.ToLower(w.Header().Get("Content-Type")), "text/event-stream")
					if responseStreaming {
						streamRewriter = newAnthropicNativeModelRewriter(publicModelID)
						w.WriteHeader(status)
						wroteHeader = true
						if flusher != nil {
							flusher.Flush()
						}
					}
					continue
				}
				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 rewriteResponse && receivedResponseStart && responseStatus >= http.StatusOK && responseStatus < http.StatusMultipleChoices {
					if responseStreaming {
						if err := writeAnthropicNativeBody(w, streamRewriter.Append(frame.GetBody()), flusher); err != nil {
							s.sendCancelRun(handle.Dispatch())
							return
						}
					} else {
						responseBody = append(responseBody, frame.GetBody()...)
					}
					continue
				}
				if !wroteHeader {
					w.Header().Set("Content-Type", "application/json")
					w.WriteHeader(http.StatusOK)
					wroteHeader = true
				}
				if err := writeAnthropicNativeBody(w, frame.GetBody(), flusher); err != nil {
					s.sendCancelRun(handle.Dispatch())
					return
				}
			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 rewriteResponse && receivedResponseStart && responseStatus >= http.StatusOK && responseStatus < http.StatusMultipleChoices {
					if responseStreaming {
						if err := writeAnthropicNativeBody(w, streamRewriter.Flush(), flusher); err != nil {
							s.sendCancelRun(handle.Dispatch())
						}
						return
					}
					if !wroteHeader {
						w.WriteHeader(responseStatus)
						wroteHeader = true
					}
					if err := writeAnthropicNativeBody(w, rewriteProviderJSONModel(responseBody, publicModelID), flusher); err != nil {
						s.sendCancelRun(handle.Dispatch())
					}
					return
				}
				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 writeAnthropicNativeBody(w http.ResponseWriter, body []byte, flusher http.Flusher) error {
	if len(body) == 0 {
		return nil
	}
	if _, err := w.Write(body); err != nil {
		return err
	}
	if flusher != nil {
		flusher.Flush()
	}
	return nil
}

type anthropicNativeModelRewriter struct {
	model        string
	pending      []byte
	messageStart bool
}

func newAnthropicNativeModelRewriter(model string) *anthropicNativeModelRewriter {
	model = strings.TrimSpace(model)
	if model == "" {
		return nil
	}
	return &anthropicNativeModelRewriter{model: model}
}

func (r *anthropicNativeModelRewriter) Append(chunk []byte) []byte {
	if r == nil || len(chunk) == 0 {
		return chunk
	}
	r.pending = append(r.pending, chunk...)
	var out bytes.Buffer
	for {
		index := bytes.IndexByte(r.pending, '\n')
		if index < 0 {
			break
		}
		line := r.pending[:index+1]
		out.Write(r.rewriteLine(line))
		r.pending = r.pending[index+1:]
	}
	return out.Bytes()
}

func (r *anthropicNativeModelRewriter) Flush() []byte {
	if r == nil || len(r.pending) == 0 {
		return nil
	}
	pending := r.pending
	r.pending = nil
	return r.rewriteLine(pending)
}

func (r *anthropicNativeModelRewriter) rewriteLine(line []byte) []byte {
	body, ending := splitLineEnding(line)
	prefix, payload, ok := bytes.Cut(body, []byte(":"))
	if !ok {
		return line
	}

	switch strings.TrimSpace(string(prefix)) {
	case "event":
		r.messageStart = strings.TrimSpace(string(payload)) == "message_start"
		return line
	case "data":
		if !r.messageStart {
			return line
		}
		r.messageStart = false
	default:
		return line
	}

	leading := len(payload) - len(bytes.TrimLeft(payload, " \t"))
	trailing := len(payload) - len(bytes.TrimRight(payload, " \t"))
	if leading+trailing >= len(payload) {
		return line
	}
	rewritten := rewriteAnthropicMessageStartModel(payload[leading:len(payload)-trailing], r.model)
	if bytes.Equal(rewritten, payload[leading:len(payload)-trailing]) {
		return line
	}
	out := make([]byte, 0, len(body)+len(rewritten)-len(payload)+len(ending))
	out = append(out, body[:len(prefix)+1+leading]...)
	out = append(out, rewritten...)
	out = append(out, payload[len(payload)-trailing:]...)
	out = append(out, ending...)
	return out
}

func rewriteAnthropicMessageStartModel(body []byte, model string) []byte {
	modelJSON, err := json.Marshal(model)
	if err != nil {
		return body
	}
	fields, _, err := scanTopLevelJSONObject(body)
	if err != nil {
		return body
	}
	for _, field := range fields {
		if field.name != "message" {
			continue
		}
		message := body[field.valueFrom:field.valueTo]
		messageFields, _, err := scanTopLevelJSONObject(message)
		if err != nil {
			return body
		}
		for _, messageField := range messageFields {
			if messageField.name != "model" {
				continue
			}
			plan, err := planTopLevelJSONPatches(message, []topLevelJSONPatch{{name: "model", value: modelJSON}})
			if err != nil {
				return body
			}
			return topLevelJSONPatchPlan{
				body: body,
				edits: []jsonByteEdit{{
					from: field.valueFrom, to: field.valueTo, replacement: plan.apply(),
				}},
				outputSize: len(body) + plan.outputSize - len(message),
			}.apply()
		}
	}
	return body
}

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)
	}
}
