package google

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

	"github.com/567-labs/instructor-go/pkg/instructor/core"
	"google.golang.org/genai"
)

func (i *InstructorGoogle) CreateChatCompletion(
	ctx context.Context,
	request GoogleRequest,
	responseType any,
) (response GoogleResponse, err error) {
	resp, err := core.ChatHandler(i, ctx, request, responseType)
	if err != nil {
		if resp == nil {
			return GoogleResponse{}, err
		}
		return *nilGoogleRespWithUsage(resp.(*GoogleResponse)), err
	}
	response = *(resp.(*GoogleResponse))
	return response, nil
}

// CreateChatCompletionUnion handles chat completion with union type extraction
// Always returns []any containing one or more variant instances
func (i *InstructorGoogle) CreateChatCompletionUnion(
	ctx context.Context,
	request GoogleRequest,
	opts core.UnionOptions,
) (result []any, response GoogleResponse, err error) {

	result, resp, err := core.ChatHandlerUnion(i, ctx, request, opts)
	if err != nil {
		if resp == nil {
			return nil, GoogleResponse{}, err
		}
		return nil, *nilGoogleRespWithUsage(resp.(*GoogleResponse)), err
	}

	response = *(resp.(*GoogleResponse))

	return result, response, nil
}

func (i *InstructorGoogle) InternalChat(ctx context.Context, request interface{}, schema *core.Schema) (string, interface{}, error) {
	req, ok := request.(GoogleRequest)
	if !ok {
		return "", nil, fmt.Errorf("invalid request type for %s client", i.Provider())
	}
	switch i.Mode() {
	case core.ModeToolCall:
		return i.chatToolCall(ctx, &req, schema, false)
	case core.ModeToolCallStrict:
		return i.chatToolCall(ctx, &req, schema, true)
	case core.ModeJSON:
		return i.chatJSON(ctx, &req, schema, false)
	case core.ModeJSONStrict:
		return i.chatJSON(ctx, &req, schema, true)
	case core.ModeJSONSchema:
		return i.chatJSONSchema(ctx, &req, schema)
	default:
		return "", nil, fmt.Errorf("mode '%s' is not supported for %s", i.Mode(), i.Provider())
	}
}

func (i *InstructorGoogle) chatToolCall(ctx context.Context, request *GoogleRequest, schema *core.Schema, strict bool) (string, *GoogleResponse, error) {
	tools := createGoogleTools(schema, strict)
	if request.GenerationConfig == nil {
		request.GenerationConfig = &genai.GenerationConfig{}
	}
	resp, err := i.Models.GenerateContent(ctx, request.Model, request.Contents, &genai.GenerateContentConfig{
		Tools:          tools,
		SafetySettings: request.SafetySettings,
	})
	if err != nil {
		return "", nil, err
	}
	googleResp := &GoogleResponse{
		Candidates:    resp.Candidates,
		UsageMetadata: resp.UsageMetadata,
	}
	var toolCalls []*genai.FunctionCall
	for _, candidate := range resp.Candidates {
		if candidate.Content != nil {
			for _, part := range candidate.Content.Parts {
				if part.FunctionCall != nil {
					toolCalls = append(toolCalls, part.FunctionCall)
				}
			}
		}
	}
	numTools := len(toolCalls)
	if numTools < 1 {
		return "", nilGoogleRespWithUsage(googleResp), errors.New("received no tool calls from model, expected at least 1")
	}

	// Helper to inject discriminator for union types
	injectDiscriminator := func(argsJSON []byte, funcName string) []byte {
		// If we have multiple functions in schema, this might be a union type
		if len(schema.Functions) > 1 {
			var argMap map[string]any
			if err := json.Unmarshal(argsJSON, &argMap); err == nil {
				// Check if discriminator already exists
				hasDiscriminator := false
				for _, v := range argMap {
					if str, ok := v.(string); ok && str == funcName {
						hasDiscriminator = true
						break
					}
				}

				// If no discriminator found, inject it as "type" field
				if !hasDiscriminator {
					argMap["type"] = funcName
					modifiedArgs, err := json.Marshal(argMap)
					if err == nil {
						return modifiedArgs
					}
				}
			}
		}
		return argsJSON
	}

	if numTools == 1 {
		argsJSON, err := json.Marshal(toolCalls[0].Args)
		if err != nil {
			return "", nilGoogleRespWithUsage(googleResp), err
		}

		argsJSON = injectDiscriminator(argsJSON, toolCalls[0].Name)
		return string(argsJSON), googleResp, nil
	}

	// Multiple tools - inject discriminator for each
	jsonArray := make([]map[string]interface{}, len(toolCalls))
	for i, toolCall := range toolCalls {
		argsJSON, err := json.Marshal(toolCall.Args)
		if err != nil {
			return "", nilGoogleRespWithUsage(googleResp), err
		}

		// Inject discriminator into arguments
		argsJSON = injectDiscriminator(argsJSON, toolCall.Name)

		var jsonObj map[string]interface{}
		err = json.Unmarshal(argsJSON, &jsonObj)
		if err != nil {
			return "", nilGoogleRespWithUsage(googleResp), err
		}
		jsonArray[i] = jsonObj
	}
	resultJSON, err := json.Marshal(jsonArray)
	if err != nil {
		return "", nilGoogleRespWithUsage(googleResp), err
	}
	return string(resultJSON), googleResp, nil
}

func (i *InstructorGoogle) chatJSON(ctx context.Context, request *GoogleRequest, schema *core.Schema, strict bool) (string, *GoogleResponse, error) {
	structName := schema.NameFromRef()
	request.Contents = prependGoogleContents(request.Contents, *createGoogleJSONMessage(schema))
	if strict {
		if request.GenerationConfig == nil {
			request.GenerationConfig = &genai.GenerationConfig{}
		}
	}
	resp, err := i.Models.GenerateContent(ctx, request.Model, request.Contents, &genai.GenerateContentConfig{
		SafetySettings: request.SafetySettings,
	})
	if err != nil {
		return "", nil, err
	}
	googleResp := &GoogleResponse{
		Candidates:    resp.Candidates,
		UsageMetadata: resp.UsageMetadata,
	}
	text := ""
	if len(resp.Candidates) > 0 && resp.Candidates[0].Content != nil {
		for _, part := range resp.Candidates[0].Content.Parts {
			if part.Text != "" {
				text = part.Text
				break
			}
		}
	}
	if strict {
		resMap := make(map[string]any)
		_ = json.Unmarshal([]byte(text), &resMap)
		cleanedText, _ := json.Marshal(resMap[structName])
		text = string(cleanedText)
	}
	return text, googleResp, nil
}

func (i *InstructorGoogle) chatJSONSchema(ctx context.Context, request *GoogleRequest, schema *core.Schema) (string, *GoogleResponse, error) {
	request.Contents = prependGoogleContents(request.Contents, *createGoogleJSONMessage(schema))
	resp, err := i.Models.GenerateContent(ctx, request.Model, request.Contents, &genai.GenerateContentConfig{
		SafetySettings: request.SafetySettings,
	})
	if err != nil {
		return "", nil, err
	}
	googleResp := &GoogleResponse{
		Candidates:    resp.Candidates,
		UsageMetadata: resp.UsageMetadata,
	}
	text := ""
	if len(resp.Candidates) > 0 && resp.Candidates[0].Content != nil {
		for _, part := range resp.Candidates[0].Content.Parts {
			if part.Text != "" {
				text = part.Text
				break
			}
		}
	}
	return text, googleResp, nil
}

func createGoogleJSONMessage(schema *core.Schema) *genai.Content {
	schemaJSON, _ := json.Marshal(schema.Schema)
	return &genai.Content{
		Parts: []*genai.Part{
			{
				Text: fmt.Sprintf("You are a helpful assistant that responds with valid JSON according to the following schema:\n\n%s\n\nRespond with valid JSON only.", string(schemaJSON)),
			},
		},
		Role: "user",
	}
}

func createGoogleTools(schema *core.Schema, strict bool) []*genai.Tool {
	// TODO: Convert schema.Schema.Properties to map[string]*genai.Schema if needed
	tool := &genai.Tool{
		FunctionDeclarations: []*genai.FunctionDeclaration{
			{
				Name:        schema.NameFromRef(),
				Description: schema.Description,
				Parameters: &genai.Schema{
					Type:       "object",
					Properties: map[string]*genai.Schema{}, // TODO: convert from schema.Schema.Properties
					Required:   []string{},                 // TODO: convert from schema.Schema.Required
				},
			},
		},
	}
	return []*genai.Tool{tool}
}

func prependGoogleContents(contents []*genai.Content, content genai.Content) []*genai.Content {
	return append([]*genai.Content{&content}, contents...)
}
