package openai

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"sort"
	"strings"
	"time"

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

type openAIChatStreamChunk struct {
	ID      string `json:"id"`
	Model   string `json:"model"`
	Choices []struct {
		Delta struct {
			Content          string `json:"content"`
			ReasoningContent string `json:"reasoning_content"`
			Reasoning        string `json:"reasoning"`
			ToolCalls        []struct {
				Index    int    `json:"index"`
				ID       string `json:"id"`
				Function struct {
					Name      string `json:"name"`
					Arguments string `json:"arguments"`
				} `json:"function"`
			} `json:"tool_calls"`
		} `json:"delta"`
		FinishReason *string `json:"finish_reason"`
	} `json:"choices"`
	Usage *openAIChatBridgeUsage `json:"usage,omitempty"`
	Error *struct {
		Type    string `json:"type"`
		Message string `json:"message"`
	} `json:"error,omitempty"`
}

type anthropicBridgeToolState struct {
	id        string
	name      string
	arguments strings.Builder
}

type anthropicBridgeStream struct {
	w          http.ResponseWriter
	model      string
	id         string
	started    bool
	stopped    bool
	nextBlock  int
	openBlock  bool
	openKind   string
	finish     *string
	usage      anthropicUsage
	tools      map[int]*anthropicBridgeToolState
	pendingSSE []byte
}

func newAnthropicBridgeStream(w http.ResponseWriter, model string) *anthropicBridgeStream {
	return &anthropicBridgeStream{w: w, model: model, tools: make(map[int]*anthropicBridgeToolState)}
}

func (s *anthropicBridgeStream) Feed(chunk []byte) error {
	if s.stopped {
		return nil
	}
	s.pendingSSE = append(s.pendingSSE, chunk...)
	s.pendingSSE = bytes.ReplaceAll(s.pendingSSE, []byte("\r\n"), []byte("\n"))
	for {
		index := bytes.Index(s.pendingSSE, []byte("\n\n"))
		if index < 0 {
			return nil
		}
		event := append([]byte(nil), s.pendingSSE[:index]...)
		s.pendingSSE = s.pendingSSE[index+2:]
		if err := s.consumeSSEEvent(event); err != nil {
			return err
		}
		if s.stopped {
			s.pendingSSE = nil
			return nil
		}
	}
}

func (s *anthropicBridgeStream) consumeSSEEvent(event []byte) error {
	var dataLines [][]byte
	for _, line := range bytes.Split(event, []byte("\n")) {
		line = bytes.TrimSpace(line)
		if len(line) == 0 || line[0] == ':' {
			continue
		}
		if bytes.HasPrefix(line, []byte("data:")) {
			dataLines = append(dataLines, bytes.TrimSpace(bytes.TrimPrefix(line, []byte("data:"))))
		}
	}
	if len(dataLines) == 0 {
		return nil
	}
	payload := bytes.Join(dataLines, []byte("\n"))
	if bytes.Equal(payload, []byte("[DONE]")) {
		return s.Finish()
	}
	var chunk openAIChatStreamChunk
	if err := json.Unmarshal(payload, &chunk); err != nil {
		return fmt.Errorf("decode Chat SSE event: %w", err)
	}
	if chunk.Error != nil {
		message := strings.TrimSpace(chunk.Error.Message)
		if message == "" {
			message = "upstream provider error"
		}
		return s.Error(chunk.Error.Type, message)
	}
	if chunk.ID != "" && s.id == "" {
		s.id = chunk.ID
	}
	if chunk.Usage != nil {
		s.usage.InputTokens = chunk.Usage.PromptTokens
		s.usage.OutputTokens = chunk.Usage.CompletionTokens
		s.usage.CacheReadInputTokens = chunk.Usage.PromptDetails.CachedTokens
	}
	if err := s.start(); err != nil {
		return err
	}
	for _, choice := range chunk.Choices {
		reasoning := choice.Delta.ReasoningContent
		if reasoning == "" {
			reasoning = choice.Delta.Reasoning
		}
		if reasoning != "" {
			if err := s.delta("thinking", reasoning); err != nil {
				return err
			}
		}
		if choice.Delta.Content != "" {
			if err := s.delta("text", choice.Delta.Content); err != nil {
				return err
			}
		}
		for _, delta := range choice.Delta.ToolCalls {
			state := s.tools[delta.Index]
			if state == nil {
				state = &anthropicBridgeToolState{}
				s.tools[delta.Index] = state
			}
			if delta.ID != "" {
				state.id = delta.ID
			}
			if delta.Function.Name != "" {
				state.name = delta.Function.Name
			}
			state.arguments.WriteString(delta.Function.Arguments)
		}
		if choice.FinishReason != nil {
			s.finish = choice.FinishReason
		}
	}
	return nil
}

func (s *anthropicBridgeStream) start() error {
	if s.started {
		return nil
	}
	s.started = true
	id := s.id
	if id == "" {
		id = "msg_iop"
	}
	return writeAnthropicSSEEvent(s.w, "message_start", map[string]any{
		"type": "message_start",
		"message": anthropicMessageResponse{
			ID: id, Type: "message", Role: "assistant", Model: s.model,
			Content: []map[string]any{}, StopReason: nil, StopSequence: nil,
			Usage: anthropicUsage{InputTokens: s.usage.InputTokens},
		},
	})
}

func (s *anthropicBridgeStream) delta(kind, value string) error {
	if !s.openBlock || s.openKind != kind {
		if err := s.closeBlock(); err != nil {
			return err
		}
		block := map[string]any{"type": kind}
		if kind == "thinking" {
			block["thinking"] = ""
			block["signature"] = ""
		} else {
			block["text"] = ""
		}
		if err := writeAnthropicSSEEvent(s.w, "content_block_start", map[string]any{
			"type": "content_block_start", "index": s.nextBlock, "content_block": block,
		}); err != nil {
			return err
		}
		s.openBlock, s.openKind = true, kind
	}
	deltaType, field := "text_delta", "text"
	if kind == "thinking" {
		deltaType, field = "thinking_delta", "thinking"
	}
	return writeAnthropicSSEEvent(s.w, "content_block_delta", map[string]any{
		"type": "content_block_delta", "index": s.nextBlock,
		"delta": map[string]any{"type": deltaType, field: value},
	})
}

func (s *anthropicBridgeStream) closeBlock() error {
	if !s.openBlock {
		return nil
	}
	if err := writeAnthropicSSEEvent(s.w, "content_block_stop", map[string]any{
		"type": "content_block_stop", "index": s.nextBlock,
	}); err != nil {
		return err
	}
	s.nextBlock++
	s.openBlock = false
	s.openKind = ""
	return nil
}

func (s *anthropicBridgeStream) emitTools() error {
	if len(s.tools) == 0 {
		return nil
	}
	indices := make([]int, 0, len(s.tools))
	for index := range s.tools {
		indices = append(indices, index)
	}
	sort.Ints(indices)
	for _, index := range indices {
		tool := s.tools[index]
		arguments := tool.arguments.String()
		if tool.id == "" || tool.name == "" || !json.Valid([]byte(arguments)) {
			return fmt.Errorf("Chat stream tool call has invalid id, name, or arguments")
		}
		if err := writeAnthropicSSEEvent(s.w, "content_block_start", map[string]any{
			"type": "content_block_start", "index": s.nextBlock,
			"content_block": map[string]any{"type": "tool_use", "id": tool.id, "name": tool.name, "input": map[string]any{}},
		}); err != nil {
			return err
		}
		if err := writeAnthropicSSEEvent(s.w, "content_block_delta", map[string]any{
			"type": "content_block_delta", "index": s.nextBlock,
			"delta": map[string]any{"type": "input_json_delta", "partial_json": arguments},
		}); err != nil {
			return err
		}
		if err := writeAnthropicSSEEvent(s.w, "content_block_stop", map[string]any{
			"type": "content_block_stop", "index": s.nextBlock,
		}); err != nil {
			return err
		}
		s.nextBlock++
	}
	return nil
}

func (s *anthropicBridgeStream) Finish() error {
	if s.stopped {
		return nil
	}
	if err := s.start(); err != nil {
		return err
	}
	if err := s.closeBlock(); err != nil {
		return err
	}
	if err := s.emitTools(); err != nil {
		return err
	}
	stopReason, err := anthropicStopReason(s.finish)
	if err != nil {
		return err
	}
	if stopReason == nil {
		fallback := "end_turn"
		stopReason = &fallback
	}
	if err := writeAnthropicSSEEvent(s.w, "message_delta", map[string]any{
		"type": "message_delta", "delta": map[string]any{"stop_reason": *stopReason, "stop_sequence": nil},
		"usage": s.usage,
	}); err != nil {
		return err
	}
	if err := writeAnthropicSSEEvent(s.w, "message_stop", map[string]any{"type": "message_stop"}); err != nil {
		return err
	}
	s.stopped = true
	return nil
}

func (s *anthropicBridgeStream) Error(errorType, message string) error {
	if s.stopped {
		return nil
	}
	if strings.TrimSpace(errorType) == "" {
		errorType = "api_error"
	}
	s.stopped = true
	return writeAnthropicSSEEvent(s.w, "error", anthropicErrorResponse{
		Type: "error", Error: errorBody{Type: errorType, Message: message},
	})
}

func writeAnthropicSSEEvent(w http.ResponseWriter, event string, payload any) error {
	encoded, err := json.Marshal(payload)
	if err != nil {
		return err
	}
	if _, err := fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event, encoded); err != nil {
		return err
	}
	return nil
}

func (s *Server) writeAnthropicChatBridgeResponse(w http.ResponseWriter, r *http.Request, handle edgeservice.ProviderTunnelResult, envelope anthropicRequestEnvelope) {
	frames := handle.Stream().Frames
	if frames == nil {
		writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider tunnel is unavailable")
		return
	}
	timer := time.NewTimer(handle.WaitTimeout())
	defer timer.Stop()
	status := http.StatusOK
	headers := make(map[string]string)
	started := false
	committed := false
	var body []byte
	var stream *anthropicBridgeStream
	if envelope.Stream {
		stream = newAnthropicBridgeStream(w, envelope.Model)
	}
	flush := func() {
		if flusher, ok := w.(http.Flusher); ok {
			flusher.Flush()
		}
	}
	commitStream := func() {
		if committed {
			return
		}
		copyAnthropicResponseHeaders(w.Header(), headers)
		w.Header().Del("Content-Length")
		w.Header().Set("Content-Type", "text/event-stream")
		w.Header().Set("Cache-Control", "no-cache")
		w.WriteHeader(status)
		committed = true
		flush()
	}

	for {
		select {
		case <-r.Context().Done():
			s.cancelRunOnHTTPGiveUp(handle.Dispatch(), r.Context().Err())
			return
		case <-timer.C:
			s.cancelRunOnHTTPGiveUp(handle.Dispatch(), errRunTimedOut)
			if !committed {
				writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider response timed out")
			} else {
				_ = stream.Error("api_error", "provider response timed out")
				flush()
			}
			return
		case frame, ok := <-frames:
			if !ok {
				if !committed {
					writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider tunnel closed before a response")
				} else {
					_ = stream.Error("api_error", "provider tunnel closed before a response")
					flush()
				}
				return
			}
			switch frame.GetKind() {
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START:
				if started {
					continue
				}
				started = true
				status = int(frame.GetStatusCode())
				if status == 0 {
					status = http.StatusOK
				}
				for key, value := range frame.GetHeaders() {
					headers[key] = value
				}
				if envelope.Stream && status < http.StatusBadRequest {
					commitStream()
				}
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY:
				if envelope.Stream && status < http.StatusBadRequest {
					commitStream()
					if err := stream.Feed(frame.GetBody()); err != nil {
						_ = stream.Error("api_error", "upstream stream could not be translated")
						flush()
						return
					}
					flush()
				} else {
					body = append(body, frame.GetBody()...)
				}
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR:
				if committed {
					_ = stream.Error("api_error", "provider tunnel failed")
					flush()
				} else {
					writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider tunnel failed")
				}
				return
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END:
				if envelope.Stream && status < http.StatusBadRequest {
					commitStream()
					if err := stream.Finish(); err != nil {
						_ = stream.Error("api_error", "upstream stream could not be translated")
					}
					flush()
					return
				}
				copyAnthropicResponseHeaders(w.Header(), headers)
				w.Header().Del("Content-Length")
				if status >= http.StatusBadRequest {
					writeJSON(w, status, convertChatErrorToAnthropic(body))
					return
				}
				converted, err := convertChatResponseToAnthropic(body, envelope.Model)
				if err != nil {
					writeAnthropicError(w, http.StatusBadGateway, "api_error", "upstream response could not be translated")
					return
				}
				writeJSON(w, http.StatusOK, converted)
				return
			case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE:
				continue
			}
		}
	}
}