package transform

import (
	"encoding/json"
	"fmt"
	"log"
	"strings"
)

// AnthropicRequest represents a standard Anthropic Messages API request.
type AnthropicRequest struct {
	Model     string             `json:"model"`
	MaxTokens int                `json:"max_tokens"`
	Messages  []AnthropicMessage `json:"messages"`
	System    interface{} `json:"system,omitempty"`
	Stream    bool               `json:"stream,omitempty"`
}

// AnthropicMessage is a single message in Anthropic format.
// Content can be a plain string or an array of content blocks.
type AnthropicMessage struct {
	Role    string      `json:"role"`
	Content interface{} `json:"content"`
}

// extractTextContent converts the message content (string or array) to a plain string.
func extractTextContent(content interface{}) string {
	if content == nil {
		return ""
	}
	switch v := content.(type) {
	case string:
		return v
	case []interface{}:
		var parts []string
		for _, block := range v {
			if b, ok := block.(map[string]interface{}); ok {
				if t, ok := b["text"].(string); ok {
					parts = append(parts, t)
				}
			}
		}
		return strings.Join(parts, "\n")
	default:
		return fmt.Sprintf("%v", v)
	}
}

// BifrostRequest is the OpenAI-compatible format that Bifrost expects.
type BifrostRequest struct {
	Model    string           `json:"model"`
	Messages []BifrostMessage `json:"messages"`
	MaxTokens int             `json:"max_tokens,omitempty"`
	Stream   bool             `json:"stream,omitempty"`
}

// BifrostMessage mirrors OpenAI chat message format.
type BifrostMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

// BifrostResponse is the OpenAI-compatible response from Bifrost.
type BifrostResponse struct {
	ID      string `json:"id"`
	Object  string `json:"object"`
	Created int64  `json:"created"`
	Model   string `json:"model"`
	Usage   struct {
		PromptTokens     int `json:"prompt_tokens"`
		CompletionTokens int `json:"completion_tokens"`
		TotalTokens      int `json:"total_tokens"`
	} `json:"usage"`
	Choices []struct {
		Message struct {
			Role      string `json:"role"`
			Content   string `json:"content"`
			Reasoning string `json:"reasoning"`
		} `json:"message"`
		FinishReason string `json:"finish_reason"`
		Index        int    `json:"index"`
	} `json:"choices"`
}

// getContent returns the message content, falling back to reasoning field
// when content is empty (DeepSeek V4 reasoning/thinking mode).
func (b BifrostResponse) getContent(idx int) string {
	if idx < len(b.Choices) {
		c := b.Choices[idx].Message.Content
		if c != "" {
			return c
		}
		return b.Choices[idx].Message.Reasoning
	}
	return ""
}

// AnthropicResponse is returned to Anthropic-format clients.
type AnthropicResponse struct {
	ID           string             `json:"id"`
	Type         string             `json:"type"`
	Role         string             `json:"role"`
	Content      []AnthropicContent `json:"content"`
	Model        string             `json:"model"`
	StopReason   string             `json:"stop_reason"`
	StopSequence *string            `json:"stop_sequence,omitempty"`
	Usage        AnthropicUsage     `json:"usage"`
}

// AnthropicUsage mirrors the Anthropic usage block.
type AnthropicUsage struct {
	InputTokens  int `json:"input_tokens"`
	OutputTokens int `json:"output_tokens"`
}

// AnthropicContent is a content block in Anthropic response format.
type AnthropicContent struct {
	Type string `json:"type"`
	Text string `json:"text"`
}

// AnthropicToBifrost converts an Anthropic Messages API request body (raw JSON)
// to Bifrost / OpenAI-compatible format.
func AnthropicToBifrost(rawBody []byte) ([]byte, error) {
	var antReq AnthropicRequest
	if err := json.Unmarshal(rawBody, &antReq); err != nil {
		return nil, fmt.Errorf("transform: parse anthropic request: %w", err)
	}

	bfMessages := make([]BifrostMessage, 0, len(antReq.Messages)+1)
	if sysContent := extractTextContent(antReq.System); sysContent != "" {
		bfMessages = append(bfMessages, BifrostMessage{Role: "system", Content: sysContent})
	}
	for _, m := range antReq.Messages {
		bfMessages = append(bfMessages, BifrostMessage{Role: m.Role, Content: extractTextContent(m.Content)})
	}

	// Auto-prefix model with provider if missing provider/model format
	model := antReq.Model
	if !strings.Contains(model, "/") {
		prefix := guessProvider(model)
		model = prefix + "/" + model
		log.Printf("[TRANSFORM] Model auto-prefixed: %s -> %s", antReq.Model, model)
	}

	bfReq := BifrostRequest{
		Model:     model,
		Messages:  bfMessages,
		MaxTokens: antReq.MaxTokens,
		Stream:    antReq.Stream,
	}
	return json.Marshal(bfReq)
}

// guessProvider maps model names to their likely provider prefix for Bifrost.
func guessProvider(model string) string {
	lower := strings.ToLower(model)
	switch {
	case strings.Contains(lower, "deepseek"):
		return "DeepSeek"
	case strings.Contains(lower, "gpt") || strings.Contains(lower, "openai"):
		return "OpenAI"
	case strings.Contains(lower, "claude") || strings.Contains(lower, "anthropic"):
		return "Anthropic"
	case strings.Contains(lower, "gemini"):
		return "Google"
	default:
		return "DeepSeek"
	}
}

// BifrostToAnthropic converts a Bifrost / OpenAI-compatible response body (raw JSON)
// back to Anthropic Messages API format.
func BifrostToAnthropic(rawBody []byte) ([]byte, error) {
	var bfResp BifrostResponse
	if err := json.Unmarshal(rawBody, &bfResp); err != nil {
		return nil, fmt.Errorf("transform: parse bifrost response: %w", err)
	}

	contents := make([]AnthropicContent, 0, len(bfResp.Choices))
	stopReason := "end_turn"
	for i := range bfResp.Choices {
		ch := bfResp.Choices[i]
		contents = append(contents, AnthropicContent{Type: "text", Text: bfResp.getContent(i)})
		if ch.FinishReason != "" {
			stopReason = ch.FinishReason
		}
	}

	antResp := AnthropicResponse{
		ID:         bfResp.ID,
		Type:       "message",
		Role:       "assistant",
		Content:    contents,
		Model:      bfResp.Model,
		StopReason: stopReason,
		Usage: AnthropicUsage{
			InputTokens:  bfResp.Usage.PromptTokens,
			OutputTokens: bfResp.Usage.CompletionTokens,
		},
	}
	return json.Marshal(antResp)
}

// ──────────────────────────────────────────────────────────────
// Streaming SSE transform: OpenAI chunks → Anthropic events
// ──────────────────────────────────────────────────────────────

// StreamTransformer converts OpenAI streaming chunks to Anthropic SSE events.
type StreamTransformer struct {
	msgID       string
	model       string
	started     bool
	blockOpened bool
	finishSent  bool
}

// NewStreamTransformer creates a new streaming transformer.
func NewStreamTransformer(msgID, model string) *StreamTransformer {
	return &StreamTransformer{msgID: msgID, model: model}
}

// TransformChunk converts a single OpenAI SSE data chunk (JSON bytes) into
// one or more Anthropic SSE event strings. Returns empty string if the
// chunk produced no output. On [DONE], emits final events.
func (s *StreamTransformer) TransformChunk(raw []byte) string {
	if string(raw) == "[DONE]" {
		return s.finish()
	}

	var chunk struct {
		ID      string `json:"id"`
		Model   string `json:"model"`
		Choices []struct {
			Delta struct {
				Content string `json:"content"`
			} `json:"delta"`
			FinishReason *string `json:"finish_reason"`
		} `json:"choices"`
	}
	if err := json.Unmarshal(raw, &chunk); err != nil {
		return ""
	}
	if len(chunk.Choices) == 0 {
		return ""
	}

	if chunk.ID != "" && chunk.Model != "" {
		s.msgID = chunk.ID
		s.model = chunk.Model
	}

	var out strings.Builder

	// message_start on first chunk
	if !s.started {
		s.started = true
		s.writeEvent(&out, "message_start", map[string]interface{}{
			"type": "message_start",
			"message": map[string]interface{}{
				"id":      s.msgID,
				"type":    "message",
				"role":    "assistant",
				"model":   s.model,
				"content": []interface{}{},
				"usage": map[string]int{
					"input_tokens":  0,
					"output_tokens": 0,
				},
			},
		})
	}

	delta := chunk.Choices[0].Delta.Content
	finish := chunk.Choices[0].FinishReason

	if delta != "" {
		if !s.blockOpened {
			s.blockOpened = true
			s.writeEvent(&out, "content_block_start", map[string]interface{}{
				"type":          "content_block_start",
				"index":         0,
				"content_block": map[string]interface{}{"type": "text", "text": ""},
			})
		}
		s.writeEvent(&out, "content_block_delta", map[string]interface{}{
			"type":  "content_block_delta",
			"index": 0,
			"delta": map[string]interface{}{"type": "text_delta", "text": delta},
		})
	}

	if finish != nil {
		out.WriteString(s.finish())
		s.finishSent = true
	}

	return out.String()
}

func (s *StreamTransformer) finish() string {
	if s.finishSent {
		return ""
	}
	var out strings.Builder
	if s.blockOpened {
		s.writeEvent(&out, "content_block_stop", map[string]interface{}{
			"type":  "content_block_stop",
			"index": 0,
		})
	}
	s.writeEvent(&out, "message_delta", map[string]interface{}{
		"type": "message_delta",
		"delta": map[string]interface{}{
			"stop_reason":   "end_turn",
			"stop_sequence": nil,
		},
		"usage": map[string]int{"output_tokens": 0},
	})
	s.writeEvent(&out, "message_stop", map[string]interface{}{
		"type": "message_stop",
	})
	s.finishSent = true
	return out.String()
}

func (s *StreamTransformer) writeEvent(out *strings.Builder, event string, data interface{}) {
	out.WriteString("event: ")
	out.WriteString(event)
	out.WriteString("\ndata: ")
	body, _ := json.Marshal(data)
	out.Write(body)
	out.WriteString("\n\n")
}