package providers

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"

	"gophergate/internal/config"
	"gophergate/internal/models"
	"github.com/go-resty/resty/v2"
)

type GeminiProvider struct {
	client *resty.Client
	config config.GeminiConfig
	apiKey string
}

func NewGeminiProvider(cfg config.GeminiConfig, apiKey string) *GeminiProvider {
	return &GeminiProvider{
		client: resty.New(),
		config: cfg,
		apiKey: apiKey,
	}
}

func (p *GeminiProvider) Name() string {
	return "gemini"
}

type GeminiRequest struct {
	Contents         []GeminiContent         `json:"contents"`
	Tools            []GeminiTool            `json:"tools,omitempty"`
	GenerationConfig *GeminiGenerationConfig `json:"generationConfig,omitempty"`
}

type GeminiTool struct {
	FunctionDeclarations []models.FunctionDef `json:"functionDeclarations"`
}

type GeminiGenerationConfig struct {
	Temperature     *float32 `json:"temperature,omitempty"`
	TopP            *float32 `json:"topP,omitempty"`
	TopK            *int     `json:"topK,omitempty"`
	MaxOutputTokens *int     `json:"maxOutputTokens,omitempty"`
	StopSequences   []string `json:"stopSequences,omitempty"`
}

type GeminiContent struct {
	Role  string       `json:"role,omitempty"`
	Parts []GeminiPart `json:"parts"`
}

type GeminiPart struct {
	Text             string                   `json:"text,omitempty"`
	InlineData       *GeminiInlineData        `json:"inlineData,omitempty"`
	FunctionCall     *GeminiFunctionCall      `json:"functionCall,omitempty"`
	FunctionResponse *GeminiFunctionResponse  `json:"functionResponse,omitempty"`
}

type GeminiInlineData struct {
	MimeType string `json:"mimeType"`
	Data     string `json:"data"`
}

type GeminiFunctionCall struct {
	Name string          `json:"name"`
	Args json.RawMessage `json:"args"`
}

type GeminiFunctionResponse struct {
	Name     string          `json:"name"`
	Response json.RawMessage `json:"response"`
}

func (p *GeminiProvider) ChatCompletion(ctx context.Context, req *models.UnifiedRequest) (*models.ChatCompletionResponse, error) {
	// Gemini mapping
	var contents []GeminiContent
	for _, msg := range req.Messages {
		role := "user"
		if msg.Role == "assistant" {
			role = "model"
		} else if msg.Role == "tool" {
			role = "function" // Function results use 'function' role in Gemini contents
		}
		
		var parts []GeminiPart
		
		// Handle tool responses
		if msg.Role == "tool" {
			text := ""
			if len(msg.Content) > 0 {
				text = msg.Content[0].Text
			}
			
			// Gemini expects functionResponse to be an object
			name := "unknown_function"
			if msg.Name != nil {
				name = *msg.Name
			}
			
			// Try to parse text as JSON if it looks like it, Gemini expects an object
			var responseObj interface{}
			if err := json.Unmarshal([]byte(text), &responseObj); err != nil {
				// If not valid JSON, wrap it in an object
				responseObj = map[string]interface{}{"result": text}
			}

			respBytes, _ := json.Marshal(responseObj)

			parts = append(parts, GeminiPart{
				FunctionResponse: &GeminiFunctionResponse{
					Name:     name,
					Response: json.RawMessage(respBytes),
				},
			})
		} else {
			for _, cp := range msg.Content {
				if cp.Type == "text" {
					parts = append(parts, GeminiPart{Text: cp.Text})
				} else if cp.Image != nil {
					base64Data, mimeType, _ := cp.Image.ToBase64()
					parts = append(parts, GeminiPart{
						InlineData: &GeminiInlineData{
							MimeType: mimeType,
							Data:     base64Data,
						},
					})
				}
			}

			// Handle assistant tool calls
			if msg.Role == "assistant" && len(msg.ToolCalls) > 0 {
				for _, tc := range msg.ToolCalls {
					parts = append(parts, GeminiPart{
						FunctionCall: &GeminiFunctionCall{
							Name: tc.Function.Name,
							Args: json.RawMessage(tc.Function.Arguments),
						},
					})
				}
			}
		}
		
		contents = append(contents, GeminiContent{
			Role:  role,
			Parts: parts,
		})
	}

	genConfig := &GeminiGenerationConfig{}
	if req.Temperature != nil {
		t := float32(*req.Temperature)
		genConfig.Temperature = &t
	}
	if req.TopP != nil {
		tp := float32(*req.TopP)
		genConfig.TopP = &tp
	}
	if req.TopK != nil {
		tk := int(*req.TopK)
		genConfig.TopK = &tk
	}
	if req.MaxTokens != nil {
		mt := int(*req.MaxTokens)
		genConfig.MaxOutputTokens = &mt
	}
	if len(req.Stop) > 0 {
		genConfig.StopSequences = req.Stop
	}

	body := GeminiRequest{
		Contents:         contents,
		GenerationConfig: genConfig,
	}

	// Map Tools
	if len(req.Tools) > 0 {
		geminiTool := GeminiTool{FunctionDeclarations: []models.FunctionDef{}}
		for _, t := range req.Tools {
			if t.Type == "function" {
				geminiTool.FunctionDeclarations = append(geminiTool.FunctionDeclarations, t.Function)
			}
		}
		body.Tools = []GeminiTool{geminiTool}
	}

	baseURL := p.config.BaseURL
	lowerModel := strings.ToLower(req.Model)
	if strings.Contains(lowerModel, "preview") || strings.Contains(lowerModel, "3.1") || strings.Contains(lowerModel, "2.0") || strings.Contains(lowerModel, "thinking") {
		// Use v1beta for preview and newer models
		if !strings.Contains(baseURL, "v1beta") {
			baseURL = strings.Replace(baseURL, "/v1", "/v1beta", 1)
		}
	}

	url := fmt.Sprintf("%s/models/%s:generateContent?key=%s", baseURL, req.Model, p.apiKey)
	fmt.Printf("[Gemini] POST %s\n", url)

	resp, err := p.client.R().
		SetContext(ctx).
		SetBody(body).
		Post(url)

	if err != nil {
		return nil, fmt.Errorf("request failed: %w", err)
	}

	if !resp.IsSuccess() {
		return nil, fmt.Errorf("Gemini API error (%d): %s", resp.StatusCode(), resp.String())
	}

	// Parse Gemini response and convert to OpenAI format
	var geminiResp struct {
		Candidates []struct {
			Content struct {
				Role  string `json:"role"`
				Parts []struct {
					Text         string              `json:"text"`
					FunctionCall *GeminiFunctionCall `json:"functionCall"`
				} `json:"parts"`
			} `json:"content"`
			FinishReason string `json:"finishReason"`
		} `json:"candidates"`
		UsageMetadata struct {
			PromptTokenCount     uint32 `json:"promptTokenCount"`
			CandidatesTokenCount uint32 `json:"candidatesTokenCount"`
			TotalTokenCount      uint32 `json:"totalTokenCount"`
		} `json:"usageMetadata"`
	}

	if err := json.Unmarshal(resp.Body(), &geminiResp); err != nil {
		return nil, fmt.Errorf("failed to parse response: %w", err)
	}

	if len(geminiResp.Candidates) == 0 {
		return nil, fmt.Errorf("no candidates in Gemini response")
	}

	content := ""
	var toolCalls []models.ToolCall
	for _, part := range geminiResp.Candidates[0].Content.Parts {
		if part.Text != "" {
			content += part.Text
		}
		if part.FunctionCall != nil {
			toolCalls = append(toolCalls, models.ToolCall{
				ID:   fmt.Sprintf("call_%s", part.FunctionCall.Name), // Gemini doesn't have call IDs
				Type: "function",
				Function: models.FunctionCall{
					Name:      part.FunctionCall.Name,
					Arguments: string(part.FunctionCall.Args),
				},
			})
		}
	}

	finishReason := strings.ToLower(geminiResp.Candidates[0].FinishReason)
	if finishReason == "stop" {
		finishReason = "stop"
	} else if len(toolCalls) > 0 {
		finishReason = "tool_calls"
	}

	openAIResp := &models.ChatCompletionResponse{
		ID:      "gemini-" + req.Model,
		Object:  "chat.completion",
		Created: 0, 
		Model:   req.Model,
		Choices: []models.ChatChoice{
			{
				Index: 0,
				Message: models.ChatMessage{
					Role:      "assistant",
					Content:   content,
					ToolCalls: toolCalls,
				},
				FinishReason: &finishReason,
			},
		},
		Usage: &models.Usage{
			PromptTokens:     geminiResp.UsageMetadata.PromptTokenCount,
			CompletionTokens: geminiResp.UsageMetadata.CandidatesTokenCount,
			TotalTokens:      geminiResp.UsageMetadata.TotalTokenCount,
		},
	}

	return openAIResp, nil
}

func (p *GeminiProvider) ChatCompletionStream(ctx context.Context, req *models.UnifiedRequest) (<-chan *models.ChatCompletionStreamResponse, error) {
	// Simplified Gemini mapping
	var contents []GeminiContent
	for _, msg := range req.Messages {
		role := "user"
		if msg.Role == "assistant" {
			role = "model"
		} else if msg.Role == "tool" {
			role = "function"
		}
		
		var parts []GeminiPart
		
		if msg.Role == "tool" {
			text := ""
			if len(msg.Content) > 0 {
				text = msg.Content[0].Text
			}
			name := "unknown"
			if msg.Name != nil {
				name = *msg.Name
			}
			
			var responseObj interface{}
			if err := json.Unmarshal([]byte(text), &responseObj); err != nil {
				responseObj = map[string]interface{}{"result": text}
			}
			respBytes, _ := json.Marshal(responseObj)

			parts = append(parts, GeminiPart{
				FunctionResponse: &GeminiFunctionResponse{
					Name:     name,
					Response: json.RawMessage(respBytes),
				},
			})
		} else {
			for _, p := range msg.Content {
				parts = append(parts, GeminiPart{Text: p.Text})
			}
			if msg.Role == "assistant" && len(msg.ToolCalls) > 0 {
				for _, tc := range msg.ToolCalls {
					parts = append(parts, GeminiPart{
						FunctionCall: &GeminiFunctionCall{
							Name: tc.Function.Name,
							Args: json.RawMessage(tc.Function.Arguments),
						},
					})
				}
			}
		}
		
		contents = append(contents, GeminiContent{
			Role:  role,
			Parts: parts,
		})
	}

	genConfig := &GeminiGenerationConfig{}
	if req.Temperature != nil {
		t := float32(*req.Temperature)
		genConfig.Temperature = &t
	}
	if req.TopP != nil {
		tp := float32(*req.TopP)
		genConfig.TopP = &tp
	}
	if req.TopK != nil {
		tk := int(*req.TopK)
		genConfig.TopK = &tk
	}
	if req.MaxTokens != nil {
		mt := int(*req.MaxTokens)
		genConfig.MaxOutputTokens = &mt
	}
	if len(req.Stop) > 0 {
		genConfig.StopSequences = req.Stop
	}

	body := GeminiRequest{
		Contents:         contents,
		GenerationConfig: genConfig,
	}

	if len(req.Tools) > 0 {
		geminiTool := GeminiTool{FunctionDeclarations: []models.FunctionDef{}}
		for _, t := range req.Tools {
			if t.Type == "function" {
				geminiTool.FunctionDeclarations = append(geminiTool.FunctionDeclarations, t.Function)
			}
		}
		body.Tools = []GeminiTool{geminiTool}
	}

	baseURL := p.config.BaseURL
	lowerModel := strings.ToLower(req.Model)
	if strings.Contains(lowerModel, "preview") || strings.Contains(lowerModel, "3.1") || strings.Contains(lowerModel, "2.0") || strings.Contains(lowerModel, "thinking") {
		// Use v1beta for preview and newer models
		if !strings.Contains(baseURL, "v1beta") {
			baseURL = strings.Replace(baseURL, "/v1", "/v1beta", 1)
		}
	}

	// Use streamGenerateContent for streaming
	url := fmt.Sprintf("%s/models/%s:streamGenerateContent?key=%s", baseURL, req.Model, p.apiKey)
	fmt.Printf("[Gemini-Stream] POST %s\n", url)

	resp, err := p.client.R().
		SetContext(ctx).
		SetBody(body).
		SetDoNotParseResponse(true).
		Post(url)

	if err != nil {
		return nil, fmt.Errorf("request failed: %w", err)
	}

	if !resp.IsSuccess() {
		return nil, fmt.Errorf("Gemini API error (%d): %s", resp.StatusCode(), resp.String())
	}

	ch := make(chan *models.ChatCompletionStreamResponse)
	
	go func() {
		defer close(ch)
		err := StreamGemini(resp.RawBody(), ch, req.Model)
		if err != nil {
			fmt.Printf("Gemini Stream error: %v\n", err)
		}
	}()

	return ch, nil
}