package stream

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"iter"
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/sirupsen/logrus"
	"google.golang.org/genai"

	"github.com/tingly-dev/tingly-box/internal/protocol"
	"github.com/tingly-dev/tingly-box/internal/protocol/nonstream"
)

// HandleGoogleToOpenAIStreamResponse processes Google streaming events and converts them to OpenAI format
func HandleGoogleToOpenAIStreamResponse(c *gin.Context, stream iter.Seq2[*genai.GenerateContentResponse, error], responseModel string) error {
	logrus.Info("Starting Google to OpenAI streaming response handler")
	defer func() {
		if r := recover(); r != nil {
			logrus.Errorf("Panic in Google to OpenAI streaming handler: %v", r)
			if c.Writer != nil {
				c.Writer.WriteHeader(http.StatusInternalServerError)
				c.Writer.Write([]byte("data: {\"error\":{\"message\":\"Internal streaming error\",\"type\":\"internal_error\"}}\n\n"))
				if flusher, ok := c.Writer.(http.Flusher); ok {
					flusher.Flush()
				}
			}
		}
		logrus.Info("Finished Google to OpenAI streaming response handler")
	}()

	// Set SSE headers
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Header("Access-Control-Allow-Origin", "*")

	flusher, ok := c.Writer.(http.Flusher)
	if !ok {
		return errors.New("Streaming not supported by this connection")
	}

	// Track streaming state
	var (
		chatID     = fmt.Sprintf("chatcmpl-%d", time.Now().Unix())
		created    = time.Now().Unix()
		toolCalls  []map[string]interface{}
		hasStarted bool
	)

	// Process the stream
	for googleResp, err := range stream {
		// Check context cancellation
		select {
		case <-c.Request.Context().Done():
			logrus.Debug("Client disconnected, stopping Google to OpenAI stream")
			return nil
		default:
		}

		if err != nil {
			// Check if it was a client cancellation
			if errors.Is(err, context.Canceled) {
				logrus.Debug("Google stream canceled by client")
				return nil
			}
			logrus.Errorf("Google stream error: %v", err)
			return nil
		}

		// Send initial chunk if not already sent
		if !hasStarted {
			chunk := map[string]interface{}{
				"id":      chatID,
				"object":  "chat.completion.chunk",
				"created": created,
				"model":   responseModel,
				"choices": []map[string]interface{}{
					{
						"index":         0,
						"delta":         map[string]interface{}{"role": "assistant"},
						"finish_reason": nil,
					},
				},
			}
			sendOpenAIStreamChunkForce(c, chunk)
			hasStarted = true
		}

		// Process candidates
		if len(googleResp.Candidates) > 0 {
			candidate := googleResp.Candidates[0]

			// Extract content
			if candidate.Content != nil {
				for _, part := range candidate.Content.Parts {
					// Handle text parts
					if part.Text != "" {
						// Send text delta
						chunk := map[string]interface{}{
							"id":      chatID,
							"object":  "chat.completion.chunk",
							"created": created,
							"model":   responseModel,
							"choices": []map[string]interface{}{
								{
									"index": 0,
									"delta": map[string]interface{}{
										"content": part.Text,
									},
									"finish_reason": nil,
								},
							},
						}
						sendOpenAIStreamChunkForce(c, chunk)
					}

					// Handle function calls
					if part.FunctionCall != nil {
						toolCall := map[string]interface{}{
							"id":   part.FunctionCall.ID,
							"type": "function",
							"function": map[string]interface{}{
								"name": part.FunctionCall.Name,
							},
						}
						// Marshal args to JSON string
						if argsBytes, err := json.Marshal(part.FunctionCall.Args); err == nil {
							toolCall["function"].(map[string]interface{})["arguments"] = string(argsBytes)
						}
						toolCalls = append(toolCalls, toolCall)

						// Send tool_calls delta
						chunk := map[string]interface{}{
							"id":      chatID,
							"object":  "chat.completion.chunk",
							"created": created,
							"model":   responseModel,
							"choices": []map[string]interface{}{
								{
									"index": 0,
									"delta": map[string]interface{}{
										"tool_calls": []map[string]interface{}{toolCall},
									},
									"finish_reason": nil,
								},
							},
						}
						sendOpenAIStreamChunkForce(c, chunk)
					}
				}
			}

			// Check for finish reason
			if candidate.FinishReason != "" {
				finishReason := nonstream.MapGoogleFinishReasonToOpenAI(candidate.FinishReason)

				// If there were tool calls, set finish reason accordingly
				if len(toolCalls) > 0 && finishReason == "stop" {
					finishReason = "tool_calls"
				}

				// Send final chunk with finish reason and usage
				chunk := map[string]interface{}{
					"id":      chatID,
					"object":  "chat.completion.chunk",
					"created": created,
					"model":   responseModel,
					"choices": []map[string]interface{}{
						{
							"index":         0,
							"delta":         map[string]interface{}{},
							"finish_reason": finishReason,
						},
					},
				}

				// Add usage if available
				if googleResp.UsageMetadata != nil {
					chunk["usage"] = map[string]interface{}{
						"prompt_tokens":     googleResp.UsageMetadata.PromptTokenCount,
						"completion_tokens": googleResp.UsageMetadata.CandidatesTokenCount,
						"total_tokens":      googleResp.UsageMetadata.TotalTokenCount,
					}
				}

				sendOpenAIStreamChunkForce(c, chunk)
				// Send final [DONE] message
				c.Writer.Write([]byte("data: [DONE]\n\n"))
				flusher.Flush()
				return nil
			}
		}
	}

	return nil
}

// HandleGoogleToAnthropicStreamResponse processes Google streaming events and converts them to Anthropic format.
// Returns UsageStat containing token usage information for tracking.
func HandleGoogleToAnthropicStreamResponse(c *gin.Context, stream iter.Seq2[*genai.GenerateContentResponse, error], responseModel string) (*protocol.TokenUsage, error) {
	logrus.Info("Starting Google to Anthropic streaming response handler")
	defer func() {
		if r := recover(); r != nil {
			logrus.Errorf("Panic in Google to Anthropic streaming handler: %v", r)
			if c.Writer != nil {
				c.Writer.WriteHeader(http.StatusInternalServerError)
				c.Writer.Write([]byte("event: error\ndata: {\"error\":{\"message\":\"Internal streaming error\",\"type\":\"internal_error\"}}\n\n"))
				if flusher, ok := c.Writer.(http.Flusher); ok {
					flusher.Flush()
				}
			}
		}
		logrus.Info("Finished Google to Anthropic streaming response handler")
	}()

	// Set SSE headers
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Header("Access-Control-Allow-Origin", "*")

	flusher, ok := c.Writer.(http.Flusher)
	if !ok {
		return protocol.ZeroTokenUsage(), errors.New("streaming not supported by this connection")
	}

	// Generate message ID for Anthropic format
	messageID := fmt.Sprintf("msg_%d", time.Now().Unix())

	// Track streaming state
	var (
		textBlockIndex = -1
		toolBlockIndex = -1
		inputTokens    int64
		outputTokens   int64
		cacheTokens    int64
	)

	// Send message_start event first
	messageStartEvent := map[string]interface{}{
		"type": "message_start",
		"message": map[string]interface{}{
			"id":            messageID,
			"type":          "message",
			"role":          "assistant",
			"content":       []interface{}{},
			"model":         responseModel,
			"stop_reason":   nil,
			"stop_sequence": nil,
			"usage": map[string]interface{}{
				"input_tokens":  0,
				"output_tokens": 0,
			},
		},
	}
	sendAnthropicStreamEvent(c, "message_start", messageStartEvent, flusher)

	// Process the stream
	for googleResp, err := range stream {
		// Check context cancellation
		select {
		case <-c.Request.Context().Done():
			logrus.Debug("Client disconnected, stopping Google to Anthropic stream")
			return protocol.NewTokenUsageWithCache(int(inputTokens), int(outputTokens), int(cacheTokens)), nil
		default:
		}

		if err != nil {
			// Check if it was a client cancellation
			if errors.Is(err, context.Canceled) {
				logrus.Debug("Google stream canceled by client")
				return protocol.NewTokenUsageWithCache(int(inputTokens), int(outputTokens), int(cacheTokens)), nil
			}
			logrus.Errorf("Google stream error: %v", err)
			errorEvent := map[string]interface{}{
				"type": "error",
				"error": map[string]interface{}{
					"message": err.Error(),
					"type":    "stream_error",
					"code":    "stream_failed",
				},
			}
			sendAnthropicStreamEvent(c, "error", errorEvent, flusher)
			return protocol.NewTokenUsageWithCache(int(inputTokens), int(outputTokens), int(cacheTokens)), err
		}

		// Process candidates
		if len(googleResp.Candidates) > 0 {
			candidate := googleResp.Candidates[0]

			// Extract content
			if candidate.Content != nil {
				for _, part := range candidate.Content.Parts {
					// Handle text parts
					if part.Text != "" {
						// Send content_block_start for text on first occurrence
						if textBlockIndex == -1 {
							textBlockIndex = 0
							contentBlockStartEvent := map[string]interface{}{
								"type":  "content_block_start",
								"index": textBlockIndex,
								"content_block": map[string]interface{}{
									"type": "text",
									"text": "",
								},
							}
							sendAnthropicStreamEvent(c, "content_block_start", contentBlockStartEvent, flusher)
						}

						// Send content_block_delta with text
						deltaEvent := map[string]interface{}{
							"type":  "content_block_delta",
							"index": textBlockIndex,
							"delta": map[string]interface{}{
								"type": "text_delta",
								"text": part.Text,
							},
						}
						sendAnthropicStreamEvent(c, "content_block_delta", deltaEvent, flusher)
					}

					// Handle function calls
					if part.FunctionCall != nil {
						toolBlockIndex++
						// Send content_block_start for tool_use
						contentBlockStartEvent := map[string]interface{}{
							"type":  "content_block_start",
							"index": toolBlockIndex,
							"content_block": map[string]interface{}{
								"type":  "tool_use",
								"id":    part.FunctionCall.ID,
								"name":  part.FunctionCall.Name,
								"input": part.FunctionCall.Args,
							},
						}
						sendAnthropicStreamEvent(c, "content_block_start", contentBlockStartEvent, flusher)

						// Send content_block_stop for this tool block
						contentBlockStopEvent := map[string]interface{}{
							"type":  "content_block_stop",
							"index": toolBlockIndex,
						}
						sendAnthropicStreamEvent(c, "content_block_stop", contentBlockStopEvent, flusher)
					}
				}
			}

			// Check for finish reason
			if candidate.FinishReason != "" {
				stopReason := nonstream.MapGoogleFinishReasonToAnthropic(candidate.FinishReason)

				// Send content_block_stop for text if applicable
				if textBlockIndex != -1 {
					contentBlockStopEvent := map[string]interface{}{
						"type":  "content_block_stop",
						"index": textBlockIndex,
					}
					sendAnthropicStreamEvent(c, "content_block_stop", contentBlockStopEvent, flusher)
				}

				// Collect usage info
				if googleResp.UsageMetadata != nil {
					inputTokens = int64(googleResp.UsageMetadata.PromptTokenCount)
					outputTokens = int64(googleResp.UsageMetadata.CandidatesTokenCount)
					cacheTokens = int64(googleResp.UsageMetadata.CachedContentTokenCount)
				}

				// Send message_delta with stop reason and usage
				messageDeltaEvent := map[string]interface{}{
					"type": "message_delta",
					"delta": map[string]interface{}{
						"stop_reason":   stopReason,
						"stop_sequence": nil,
					},
					"usage": map[string]interface{}{
						"output_tokens": outputTokens,
					},
				}
				sendAnthropicStreamEvent(c, "message_delta", messageDeltaEvent, flusher)

				// Send message_stop
				messageStopEvent := map[string]interface{}{
					"type": "message_stop",
				}
				sendAnthropicStreamEvent(c, "message_stop", messageStopEvent, flusher)
				return protocol.NewTokenUsageWithCache(int(inputTokens), int(outputTokens), int(cacheTokens)), nil
			}
		}

		// Track usage
		if googleResp.UsageMetadata != nil {
			inputTokens = int64(googleResp.UsageMetadata.PromptTokenCount)
			outputTokens = int64(googleResp.UsageMetadata.CandidatesTokenCount)
			cacheTokens = int64(googleResp.UsageMetadata.CachedContentTokenCount)
		}
	}

	return protocol.NewTokenUsageWithCache(int(inputTokens), int(outputTokens), int(cacheTokens)), nil
}

// HandleGoogleToAnthropicBetaStreamResponse processes Google streaming events and converts them to Anthropic beta format.
// Returns UsageStat containing token usage information for tracking.
func HandleGoogleToAnthropicBetaStreamResponse(c *gin.Context, stream iter.Seq2[*genai.GenerateContentResponse, error], responseModel string) (*protocol.TokenUsage, error) {
	logrus.Info("Starting Google to Anthropic beta streaming response handler")
	defer func() {
		if r := recover(); r != nil {
			logrus.Errorf("Panic in Google to Anthropic beta streaming handler: %v", r)
			if c.Writer != nil {
				c.Writer.WriteHeader(http.StatusInternalServerError)
				c.Writer.Write([]byte("event: error\ndata: {\"error\":{\"message\":\"Internal streaming error\",\"type\":\"internal_error\"}}\n\n"))
				if flusher, ok := c.Writer.(http.Flusher); ok {
					flusher.Flush()
				}
			}
		}
		logrus.Info("Finished Google to Anthropic beta streaming response handler")
	}()

	// Set SSE headers
	c.Header("Content-Type", "text/event-stream")
	c.Header("Cache-Control", "no-cache")
	c.Header("Connection", "keep-alive")
	c.Header("Access-Control-Allow-Origin", "*")

	flusher, ok := c.Writer.(http.Flusher)
	if !ok {
		return protocol.ZeroTokenUsage(), errors.New("streaming not supported by this connection")
	}

	// Generate message ID for Anthropic beta format
	messageID := fmt.Sprintf("msg_%d", time.Now().Unix())

	// Track streaming state
	var (
		textBlockIndex = -1
		toolBlockIndex = -1
		inputTokens    int64
		outputTokens   int64
		cacheTokens    int64
	)

	// Send message_start event first
	messageStartEvent := map[string]interface{}{
		"type": eventTypeMessageStart,
		"message": map[string]interface{}{
			"id":            messageID,
			"type":          "message",
			"role":          "assistant",
			"content":       []interface{}{},
			"model":         responseModel,
			"stop_reason":   nil,
			"stop_sequence": nil,
			"usage": map[string]interface{}{
				"input_tokens":  0,
				"output_tokens": 0,
			},
		},
	}
	sendAnthropicStreamEvent(c, eventTypeMessageStart, messageStartEvent, flusher)

	// Process the stream
	for googleResp, err := range stream {
		// Check context cancellation
		select {
		case <-c.Request.Context().Done():
			logrus.Debug("Client disconnected, stopping Google to Anthropic beta stream")
			return protocol.NewTokenUsageWithCache(int(inputTokens), int(outputTokens), int(cacheTokens)), nil
		default:
		}

		if err != nil {
			// Check if it was a client cancellation
			if errors.Is(err, context.Canceled) {
				logrus.Debug("Google stream canceled by client")
				return protocol.NewTokenUsageWithCache(int(inputTokens), int(outputTokens), int(cacheTokens)), nil
			}
			logrus.Errorf("Google stream error: %v", err)
			errorEvent := map[string]interface{}{
				"type": "error",
				"error": map[string]interface{}{
					"message": err.Error(),
					"type":    "stream_error",
					"code":    "stream_failed",
				},
			}
			sendAnthropicStreamEvent(c, "error", errorEvent, flusher)
			return protocol.NewTokenUsageWithCache(int(inputTokens), int(outputTokens), int(cacheTokens)), err
		}

		// Process candidates
		if len(googleResp.Candidates) > 0 {
			candidate := googleResp.Candidates[0]

			// Extract content
			if candidate.Content != nil {
				for _, part := range candidate.Content.Parts {
					// Handle text parts
					if part.Text != "" {
						// Send content_block_start for text on first occurrence
						if textBlockIndex == -1 {
							textBlockIndex = 0
							contentBlockStartEvent := map[string]interface{}{
								"type":  eventTypeContentBlockStart,
								"index": textBlockIndex,
								"content_block": map[string]interface{}{
									"type": "text",
									"text": "",
								},
							}
							sendAnthropicStreamEvent(c, eventTypeContentBlockStart, contentBlockStartEvent, flusher)
						}

						// Send content_block_delta with text
						deltaEvent := map[string]interface{}{
							"type":  eventTypeContentBlockDelta,
							"index": textBlockIndex,
							"delta": map[string]interface{}{
								"type": "text_delta",
								"text": part.Text,
							},
						}
						sendAnthropicStreamEvent(c, eventTypeContentBlockDelta, deltaEvent, flusher)
					}

					// Handle function calls
					if part.FunctionCall != nil {
						toolBlockIndex++
						// Send content_block_start for tool_use
						contentBlockStartEvent := map[string]interface{}{
							"type":  eventTypeContentBlockStart,
							"index": toolBlockIndex,
							"content_block": map[string]interface{}{
								"type":  "tool_use",
								"id":    part.FunctionCall.ID,
								"name":  part.FunctionCall.Name,
								"input": part.FunctionCall.Args,
							},
						}
						sendAnthropicStreamEvent(c, eventTypeContentBlockStart, contentBlockStartEvent, flusher)

						// Send content_block_stop for this tool block
						contentBlockStopEvent := map[string]interface{}{
							"type":  eventTypeContentBlockStop,
							"index": toolBlockIndex,
						}
						sendAnthropicStreamEvent(c, eventTypeContentBlockStop, contentBlockStopEvent, flusher)
					}
				}
			}

			// Check for finish reason
			if candidate.FinishReason != "" {
				stopReason := nonstream.MapGoogleFinishReasonToAnthropicBeta(candidate.FinishReason)

				// Send content_block_stop for text if applicable
				if textBlockIndex != -1 {
					contentBlockStopEvent := map[string]interface{}{
						"type":  eventTypeContentBlockStop,
						"index": textBlockIndex,
					}
					sendAnthropicStreamEvent(c, eventTypeContentBlockStop, contentBlockStopEvent, flusher)
				}

				// Collect usage info
				if googleResp.UsageMetadata != nil {
					inputTokens = int64(googleResp.UsageMetadata.PromptTokenCount)
					outputTokens = int64(googleResp.UsageMetadata.CandidatesTokenCount)
					cacheTokens = int64(googleResp.UsageMetadata.CachedContentTokenCount)
				}

				// Send message_delta with stop reason and usage
				messageDeltaEvent := map[string]interface{}{
					"type": eventTypeMessageDelta,
					"delta": map[string]interface{}{
						"stop_reason":   string(stopReason),
						"stop_sequence": "",
					},
					"usage": map[string]interface{}{
						"output_tokens": outputTokens,
					},
				}
				sendAnthropicStreamEvent(c, eventTypeMessageDelta, messageDeltaEvent, flusher)

				// Send message_stop
				messageStopEvent := map[string]interface{}{
					"type": eventTypeMessageStop,
					"message": map[string]interface{}{
						"id":            messageID,
						"type":          "message",
						"role":          "assistant",
						"content":       []interface{}{},
						"model":         responseModel,
						"stop_reason":   string(stopReason),
						"stop_sequence": "",
						"usage": map[string]interface{}{
							"output_tokens": outputTokens,
						},
					},
				}
				sendAnthropicStreamEvent(c, eventTypeMessageStop, messageStopEvent, flusher)

				// Send final simple data with type (without event, aka empty)
				c.SSEvent("", map[string]interface{}{"type": eventTypeMessageStop})
				flusher.Flush()
				return protocol.NewTokenUsageWithCache(int(inputTokens), int(outputTokens), int(cacheTokens)), nil
			}
		}

		// Track usage
		if googleResp.UsageMetadata != nil {
			inputTokens = int64(googleResp.UsageMetadata.PromptTokenCount)
			outputTokens = int64(googleResp.UsageMetadata.CandidatesTokenCount)
			cacheTokens = int64(googleResp.UsageMetadata.CachedContentTokenCount)
		}
	}

	return protocol.NewTokenUsageWithCache(int(inputTokens), int(outputTokens), int(cacheTokens)), nil
}
