package proxy

import (
	"bufio"
	"bytes"
	"encoding/json"
	"strings"
)

func assembleSSEJSON(body []byte) ([]byte, bool) {
	payloads := sseDataPayloads(body)
	if len(payloads) == 0 {
		return nil, false
	}
	if assembled, ok := assembleOpenAICompletionsSSE(payloads); ok {
		return assembled, true
	}
	if assembled, ok := assembleOpenAIChatSSE(payloads); ok {
		return assembled, true
	}
	if assembled, ok := assembleOpenAIResponsesSSE(payloads); ok {
		return assembled, true
	}
	if assembled, ok := assembleAnthropicSSE(payloads); ok {
		return assembled, true
	}
	if assembled, ok := assembleGeminiSSE(payloads); ok {
		return assembled, true
	}
	return nil, false
}

func sseDataPayloads(body []byte) []string {
	payloads := make([]string, 0)
	scanner := bufio.NewScanner(bytes.NewReader(body))
	scanner.Buffer(make([]byte, 0, 64*1024), len(body)+1)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
		if payload == "" || payload == "[DONE]" {
			continue
		}
		payloads = append(payloads, payload)
	}
	return payloads
}

type openAICompletionChunk struct {
	ID      string          `json:"id,omitempty"`
	Object  string          `json:"object,omitempty"`
	Created int64           `json:"created,omitempty"`
	Model   string          `json:"model,omitempty"`
	Usage   json.RawMessage `json:"usage"`
	Choices []struct {
		Index        int     `json:"index"`
		Text         *string `json:"text"`
		FinishReason *string `json:"finish_reason"`
	} `json:"choices"`
}

type openAICompletionChoice struct {
	index        int
	text         strings.Builder
	finishReason string
}

func assembleOpenAICompletionsSSE(payloads []string) ([]byte, bool) {
	choices := map[int]*openAICompletionChoice{}
	order := make([]int, 0, 1)
	var id, object, model string
	var created int64
	var usage json.RawMessage
	matched := false

	for _, payload := range payloads {
		var chunk openAICompletionChunk
		if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
			return nil, false
		}
		if len(chunk.Choices) == 0 && len(chunk.Usage) == 0 {
			return nil, false
		}
		for _, choice := range chunk.Choices {
			if choice.Text == nil {
				return nil, false
			}
		}
		matched = true
		if id == "" {
			id = chunk.ID
		}
		if object == "" {
			object = chunk.Object
		}
		if created == 0 {
			created = chunk.Created
		}
		if model == "" {
			model = chunk.Model
		}
		if len(chunk.Usage) > 0 && string(chunk.Usage) != "null" {
			usage = chunk.Usage
		}

		for _, choice := range chunk.Choices {
			assembled := choices[choice.Index]
			if assembled == nil {
				assembled = &openAICompletionChoice{index: choice.Index}
				choices[choice.Index] = assembled
				order = append(order, choice.Index)
			}
			assembled.text.WriteString(*choice.Text)
			if choice.FinishReason != nil {
				assembled.finishReason = *choice.FinishReason
			}
		}
	}
	if !matched {
		return nil, false
	}

	assembled := struct {
		ID      string          `json:"id,omitempty"`
		Object  string          `json:"object,omitempty"`
		Created int64           `json:"created,omitempty"`
		Model   string          `json:"model,omitempty"`
		Usage   json.RawMessage `json:"usage,omitempty"`
		Choices []struct {
			Index        int    `json:"index"`
			Text         string `json:"text"`
			FinishReason string `json:"finish_reason,omitempty"`
		} `json:"choices"`
	}{ID: id, Object: object, Created: created, Model: model, Usage: usage}
	for _, index := range order {
		choice := choices[index]
		assembled.Choices = append(assembled.Choices, struct {
			Index        int    `json:"index"`
			Text         string `json:"text"`
			FinishReason string `json:"finish_reason,omitempty"`
		}{Index: choice.index, Text: choice.text.String(), FinishReason: choice.finishReason})
	}

	data, err := json.Marshal(assembled)
	return data, err == nil
}

type openAIChatChunk struct {
	ID                string          `json:"id,omitempty"`
	Object            string          `json:"object,omitempty"`
	Created           int64           `json:"created,omitempty"`
	Model             string          `json:"model,omitempty"`
	SystemFingerprint string          `json:"system_fingerprint,omitempty"`
	Usage             json.RawMessage `json:"usage"`
	Choices           []struct {
		Index int `json:"index"`
		Delta struct {
			Role             string `json:"role,omitempty"`
			Content          string `json:"content,omitempty"`
			ReasoningContent string `json:"reasoning_content,omitempty"`
			ToolCalls        []struct {
				Index    int    `json:"index"`
				ID       string `json:"id,omitempty"`
				Type     string `json:"type,omitempty"`
				Function struct {
					Name      string `json:"name,omitempty"`
					Arguments string `json:"arguments,omitempty"`
				} `json:"function,omitempty"`
			} `json:"tool_calls,omitempty"`
		} `json:"delta"`
		FinishReason       *string `json:"finish_reason"`
		NativeFinishReason *string `json:"native_finish_reason"`
	} `json:"choices"`
}

type openAIChatChoice struct {
	index        int
	role         string
	content      strings.Builder
	reasoning    strings.Builder
	toolCalls    map[int]*openAIToolCall
	toolOrder    []int
	finishReason string
	nativeFinish string
}

type openAIToolCall struct {
	id        string
	callType  string
	name      string
	arguments strings.Builder
}

type openAIChatResponse struct {
	ID                string                 `json:"id,omitempty"`
	Object            string                 `json:"object,omitempty"`
	Created           int64                  `json:"created,omitempty"`
	Model             string                 `json:"model,omitempty"`
	SystemFingerprint string                 `json:"system_fingerprint,omitempty"`
	Usage             json.RawMessage        `json:"usage,omitempty"`
	Choices           []openAIResponseChoice `json:"choices"`
}

type openAIResponseChoice struct {
	Index              int                   `json:"index"`
	Message            openAIResponseMessage `json:"message"`
	FinishReason       string                `json:"finish_reason,omitempty"`
	NativeFinishReason string                `json:"native_finish_reason,omitempty"`
}

type openAIResponseMessage struct {
	Role             string               `json:"role,omitempty"`
	Content          string               `json:"content"`
	ReasoningContent string               `json:"reasoning_content,omitempty"`
	ToolCalls        []openAIResponseTool `json:"tool_calls,omitempty"`
}

type openAIResponseTool struct {
	ID       string                     `json:"id,omitempty"`
	Type     string                     `json:"type,omitempty"`
	Function openAIResponseToolFunction `json:"function"`
}

type openAIResponseToolFunction struct {
	Name      string `json:"name,omitempty"`
	Arguments string `json:"arguments"`
}

func assembleOpenAIChatSSE(payloads []string) ([]byte, bool) {
	choices := map[int]*openAIChatChoice{}
	order := make([]int, 0, 1)
	var id, object string
	var created int64
	var model string
	var systemFingerprint string
	var usage json.RawMessage
	matched := false

	for _, payload := range payloads {
		var chunk openAIChatChunk
		if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
			return nil, false
		}
		if len(chunk.Choices) == 0 && len(chunk.Usage) == 0 {
			return nil, false
		}
		matched = true
		if id == "" {
			id = chunk.ID
		}
		if object == "" {
			object = strings.TrimSuffix(chunk.Object, ".chunk")
		}
		if created == 0 {
			created = chunk.Created
		}
		if model == "" {
			model = chunk.Model
		}
		if systemFingerprint == "" {
			systemFingerprint = chunk.SystemFingerprint
		}
		if len(chunk.Usage) > 0 && string(chunk.Usage) != "null" {
			usage = chunk.Usage
		}

		for _, choice := range chunk.Choices {
			assembled := choices[choice.Index]
			if assembled == nil {
				assembled = &openAIChatChoice{index: choice.Index}
				choices[choice.Index] = assembled
				order = append(order, choice.Index)
			}
			if choice.Delta.Role != "" {
				assembled.role = choice.Delta.Role
			}
			if choice.Delta.Content != "" {
				assembled.content.WriteString(choice.Delta.Content)
			}
			if choice.Delta.ReasoningContent != "" {
				assembled.reasoning.WriteString(choice.Delta.ReasoningContent)
			}
			for _, toolCall := range choice.Delta.ToolCalls {
				if assembled.toolCalls == nil {
					assembled.toolCalls = map[int]*openAIToolCall{}
				}
				assembledTool := assembled.toolCalls[toolCall.Index]
				if assembledTool == nil {
					assembledTool = &openAIToolCall{}
					assembled.toolCalls[toolCall.Index] = assembledTool
					assembled.toolOrder = append(assembled.toolOrder, toolCall.Index)
				}
				if toolCall.ID != "" {
					assembledTool.id = toolCall.ID
				}
				if toolCall.Type != "" {
					assembledTool.callType = toolCall.Type
				}
				if toolCall.Function.Name != "" {
					assembledTool.name = toolCall.Function.Name
				}
				if toolCall.Function.Arguments != "" {
					assembledTool.arguments.WriteString(toolCall.Function.Arguments)
				}
			}
			if choice.FinishReason != nil {
				assembled.finishReason = *choice.FinishReason
			}
			if choice.NativeFinishReason != nil {
				assembled.nativeFinish = *choice.NativeFinishReason
			}
		}
	}
	if !matched {
		return nil, false
	}

	assembled := openAIChatResponse{ID: id, Object: object, Created: created, Model: model, SystemFingerprint: systemFingerprint, Usage: usage}
	for _, index := range order {
		choice := choices[index]
		out := openAIResponseChoice{Index: choice.index, FinishReason: choice.finishReason, NativeFinishReason: choice.nativeFinish}
		out.Message.Role = choice.role
		out.Message.Content = choice.content.String()
		out.Message.ReasoningContent = choice.reasoning.String()
		for _, toolIndex := range choice.toolOrder {
			toolCall := choice.toolCalls[toolIndex]
			outTool := openAIResponseTool{ID: toolCall.id, Type: toolCall.callType}
			outTool.Function.Name = toolCall.name
			outTool.Function.Arguments = toolCall.arguments.String()
			out.Message.ToolCalls = append(out.Message.ToolCalls, outTool)
		}
		assembled.Choices = append(assembled.Choices, out)
	}

	data, err := json.Marshal(assembled)
	return data, err == nil
}

type openAIResponsesEvent struct {
	Type     string          `json:"type"`
	Response json.RawMessage `json:"response"`
}

func assembleOpenAIResponsesSSE(payloads []string) ([]byte, bool) {
	var completed json.RawMessage
	matched := false
	for _, payload := range payloads {
		var event openAIResponsesEvent
		if err := json.Unmarshal([]byte(payload), &event); err != nil {
			return nil, false
		}
		if !strings.HasPrefix(event.Type, "response.") {
			return nil, false
		}
		matched = true
		if event.Type == "response.completed" && len(event.Response) > 0 {
			completed = append(json.RawMessage(nil), event.Response...)
		}
	}
	if !matched || len(completed) == 0 {
		return nil, false
	}
	return completed, true
}

type anthropicEvent struct {
	Type    string `json:"type"`
	Message *struct {
		ID           string `json:"id"`
		Type         string `json:"type"`
		Role         string `json:"role"`
		Model        string `json:"model,omitempty"`
		StopReason   string `json:"stop_reason"`
		StopSequence string `json:"stop_sequence"`
		Usage        struct {
			InputTokens  int `json:"input_tokens"`
			OutputTokens int `json:"output_tokens"`
		} `json:"usage"`
	} `json:"message"`
	Index        int `json:"index"`
	ContentBlock *struct {
		Type string `json:"type"`
		Text string `json:"text"`
		ID   string `json:"id"`
		Name string `json:"name"`
	} `json:"content_block"`
	Delta *struct {
		StopReason   string `json:"stop_reason"`
		StopSequence string `json:"stop_sequence"`
		Text         string `json:"text"`
		PartialJSON  string `json:"partial_json"`
	} `json:"delta"`
	Usage *struct {
		OutputTokens int `json:"output_tokens"`
	} `json:"usage"`
}

type anthropicContentBlock struct {
	Type  string         `json:"type"`
	Text  string         `json:"text,omitempty"`
	ID    string         `json:"id,omitempty"`
	Name  string         `json:"name,omitempty"`
	Input map[string]any `json:"input,omitempty"`

	partialJSON strings.Builder
}

func assembleAnthropicSSE(payloads []string) ([]byte, bool) {
	var assembled struct {
		ID           string                  `json:"id"`
		Type         string                  `json:"type"`
		Role         string                  `json:"role"`
		Content      []anthropicContentBlock `json:"content"`
		Model        string                  `json:"model,omitempty"`
		StopReason   string                  `json:"stop_reason,omitempty"`
		StopSequence string                  `json:"stop_sequence,omitempty"`
		Usage        struct {
			InputTokens  int `json:"input_tokens"`
			OutputTokens int `json:"output_tokens"`
		} `json:"usage"`
	}
	contents := map[int]*anthropicContentBlock{}
	order := make([]int, 0, 1)
	matched := false

	for _, payload := range payloads {
		var event anthropicEvent
		if err := json.Unmarshal([]byte(payload), &event); err != nil {
			return nil, false
		}
		if !strings.HasPrefix(event.Type, "message_") && !strings.HasPrefix(event.Type, "content_block_") {
			return nil, false
		}
		matched = true

		if event.Message != nil {
			assembled.ID = event.Message.ID
			assembled.Type = event.Message.Type
			assembled.Role = event.Message.Role
			assembled.Model = event.Message.Model
			assembled.StopReason = event.Message.StopReason
			assembled.StopSequence = event.Message.StopSequence
			assembled.Usage.InputTokens = event.Message.Usage.InputTokens
			assembled.Usage.OutputTokens = event.Message.Usage.OutputTokens
		}
		if event.ContentBlock != nil {
			block := contents[event.Index]
			if block == nil {
				block = &anthropicContentBlock{Type: event.ContentBlock.Type, ID: event.ContentBlock.ID, Name: event.ContentBlock.Name}
				contents[event.Index] = block
				order = append(order, event.Index)
			}
			if event.ContentBlock.ID != "" {
				block.ID = event.ContentBlock.ID
			}
			if event.ContentBlock.Name != "" {
				block.Name = event.ContentBlock.Name
			}
			block.Text += event.ContentBlock.Text
		}
		if event.Delta != nil {
			if event.Delta.Text != "" {
				block := contents[event.Index]
				if block == nil {
					block = &anthropicContentBlock{Type: "text"}
					contents[event.Index] = block
					order = append(order, event.Index)
				}
				block.Text += event.Delta.Text
			}
			if event.Delta.PartialJSON != "" {
				block := contents[event.Index]
				if block == nil {
					block = &anthropicContentBlock{Type: "tool_use"}
					contents[event.Index] = block
					order = append(order, event.Index)
				}
				block.partialJSON.WriteString(event.Delta.PartialJSON)
			}
			if event.Delta.StopReason != "" {
				assembled.StopReason = event.Delta.StopReason
			}
			if event.Delta.StopSequence != "" {
				assembled.StopSequence = event.Delta.StopSequence
			}
		}
		if event.Usage != nil {
			assembled.Usage.OutputTokens = event.Usage.OutputTokens
		}
	}
	if !matched || assembled.Type == "" {
		return nil, false
	}
	for _, index := range order {
		block := contents[index]
		if block.partialJSON.Len() > 0 {
			var input map[string]any
			if err := json.Unmarshal([]byte(block.partialJSON.String()), &input); err != nil {
				return nil, false
			}
			block.Input = input
		}
		assembled.Content = append(assembled.Content, *block)
	}

	data, err := json.Marshal(assembled)
	return data, err == nil
}

type geminiChunk struct {
	Raw        map[string]json.RawMessage `json:"-"`
	Candidates []struct {
		Raw     map[string]json.RawMessage `json:"-"`
		Content struct {
			Parts []struct {
				Text string `json:"text"`
			} `json:"parts"`
			Role string `json:"role"`
		} `json:"content"`
		FinishReason string `json:"finishReason"`
		Index        int    `json:"index"`
	} `json:"candidates"`
	UsageMetadata json.RawMessage `json:"usageMetadata"`
}

func (g *geminiChunk) UnmarshalJSON(data []byte) error {
	type alias geminiChunk
	var decoded alias
	if err := json.Unmarshal(data, &decoded); err != nil {
		return err
	}
	var raw map[string]json.RawMessage
	if err := json.Unmarshal(data, &raw); err != nil {
		return err
	}
	*g = geminiChunk(decoded)
	g.Raw = raw
	if candidatesRaw, ok := raw["candidates"]; ok {
		var rawCandidates []map[string]json.RawMessage
		if err := json.Unmarshal(candidatesRaw, &rawCandidates); err != nil {
			return err
		}
		for i := range g.Candidates {
			if i < len(rawCandidates) {
				g.Candidates[i].Raw = rawCandidates[i]
			}
		}
	}
	return nil
}

type geminiCandidate struct {
	raw   map[string]json.RawMessage
	index int
	role  string
	parts []*geminiPart
}

type geminiPart struct {
	raw  map[string]json.RawMessage
	text strings.Builder
	kind string
}

func assembleGeminiSSE(payloads []string) ([]byte, bool) {
	candidates := map[int]*geminiCandidate{}
	order := make([]int, 0, 1)
	var usage json.RawMessage
	matched := false

	for _, payload := range payloads {
		var chunk geminiChunk
		if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
			return nil, false
		}
		if len(chunk.Candidates) == 0 {
			return nil, false
		}
		matched = true
		for _, candidate := range chunk.Candidates {
			assembled := candidates[candidate.Index]
			if assembled == nil {
				assembled = &geminiCandidate{index: candidate.Index}
				candidates[candidate.Index] = assembled
				order = append(order, candidate.Index)
			}
			if candidate.Raw != nil {
				assembled.raw = cloneRawMap(candidate.Raw)
			}
			if candidate.Content.Role != "" {
				assembled.role = candidate.Content.Role
			}
			if len(candidate.Content.Parts) > 0 {
				var contentRaw struct {
					Parts []map[string]json.RawMessage `json:"parts"`
				}
				if rawContent, ok := candidate.Raw["content"]; ok {
					_ = json.Unmarshal(rawContent, &contentRaw)
				}
				mergeByIndex := len(assembled.parts) == len(candidate.Content.Parts)
				if mergeByIndex {
					for i := range candidate.Content.Parts {
						if i >= len(contentRaw.Parts) || assembled.parts[i].kind != geminiPartKind(contentRaw.Parts[i]) {
							mergeByIndex = false
							break
						}
					}
				}
				for i, part := range candidate.Content.Parts {
					target := i
					if !mergeByIndex {
						target = len(assembled.parts)
						assembled.parts = append(assembled.parts, &geminiPart{})
					}
					assembled.parts[target].text.WriteString(part.Text)
					if i < len(contentRaw.Parts) {
						assembled.parts[target].kind = geminiPartKind(contentRaw.Parts[i])
						assembled.parts[target].raw = mergeRawMap(assembled.parts[target].raw, contentRaw.Parts[i])
					}
				}
			}
		}
		if len(chunk.UsageMetadata) > 0 {
			usage = chunk.UsageMetadata
		}
	}
	if !matched {
		return nil, false
	}

	assembled := map[string]any{"candidates": make([]any, 0, len(order))}
	for _, index := range order {
		candidate := candidates[index]
		candidateRaw := cloneRawMap(candidate.raw)
		candidateRaw["index"] = mustJSON(candidate.index)
		parts := make([]any, 0, len(candidate.parts))
		for _, part := range candidate.parts {
			partRaw := cloneRawMap(part.raw)
			if part.text.Len() > 0 || len(partRaw) == 0 {
				partRaw["text"] = mustJSON(part.text.String())
			}
			parts = append(parts, rawMapToMap(partRaw))
		}
		contentRaw := map[string]any{"role": candidate.role, "parts": parts}
		candidateRaw["content"] = mustJSON(contentRaw)
		assembled["candidates"] = append(assembled["candidates"].([]any), rawMapToMap(candidateRaw))
	}
	if len(usage) > 0 {
		assembled["usageMetadata"] = usage
	}

	data, err := json.Marshal(assembled)
	return data, err == nil
}

func geminiPartKind(part map[string]json.RawMessage) string {
	for _, key := range []string{"functionCall", "functionResponse", "inlineData", "fileData", "executableCode", "codeExecutionResult", "text"} {
		if _, ok := part[key]; ok {
			return key
		}
	}
	return "unknown"
}

func mergeRawMap(dst, src map[string]json.RawMessage) map[string]json.RawMessage {
	if dst == nil {
		dst = make(map[string]json.RawMessage, len(src))
	}
	for key, value := range src {
		if key == "text" {
			continue
		}
		dst[key] = append(json.RawMessage(nil), value...)
	}
	return dst
}

func cloneRawMap(in map[string]json.RawMessage) map[string]json.RawMessage {
	out := make(map[string]json.RawMessage, len(in))
	for k, v := range in {
		out[k] = append(json.RawMessage(nil), v...)
	}
	return out
}

func mustJSON(v any) json.RawMessage {
	data, err := json.Marshal(v)
	if err != nil {
		return nil
	}
	return data
}

func rawMapToMap(in map[string]json.RawMessage) map[string]any {
	out := make(map[string]any, len(in))
	for k, v := range in {
		var decoded any
		if err := json.Unmarshal(v, &decoded); err != nil {
			out[k] = string(v)
			continue
		}
		out[k] = decoded
	}
	return out
}