package openai

import (
	"cmp"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"reflect"
	"slices"
	"strings"

	"charm.land/fantasy"
	"charm.land/fantasy/object"
	"charm.land/fantasy/schema"
	"github.com/charmbracelet/openai-go"
	"github.com/charmbracelet/openai-go/packages/param"
	"github.com/charmbracelet/openai-go/shared"
	"github.com/google/uuid"
)

type languageModel struct {
	provider                   string
	modelID                    string
	client                     openai.Client
	objectMode                 fantasy.ObjectMode
	prepareCallFunc            LanguageModelPrepareCallFunc
	mapFinishReasonFunc        LanguageModelMapFinishReasonFunc
	extraContentFunc           LanguageModelExtraContentFunc
	usageFunc                  LanguageModelUsageFunc
	streamUsageFunc            LanguageModelStreamUsageFunc
	streamExtraFunc            LanguageModelStreamExtraFunc
	streamProviderMetadataFunc LanguageModelStreamProviderMetadataFunc
	toPromptFunc               LanguageModelToPromptFunc
}

// LanguageModelOption is a function that configures a languageModel.
type LanguageModelOption = func(*languageModel)

// WithLanguageModelPrepareCallFunc sets the prepare call function for the language model.
func WithLanguageModelPrepareCallFunc(fn LanguageModelPrepareCallFunc) LanguageModelOption {
	return func(l *languageModel) {
		l.prepareCallFunc = fn
	}
}

// WithLanguageModelMapFinishReasonFunc sets the map finish reason function for the language model.
func WithLanguageModelMapFinishReasonFunc(fn LanguageModelMapFinishReasonFunc) LanguageModelOption {
	return func(l *languageModel) {
		l.mapFinishReasonFunc = fn
	}
}

// WithLanguageModelExtraContentFunc sets the extra content function for the language model.
func WithLanguageModelExtraContentFunc(fn LanguageModelExtraContentFunc) LanguageModelOption {
	return func(l *languageModel) {
		l.extraContentFunc = fn
	}
}

// WithLanguageModelStreamExtraFunc sets the stream extra function for the language model.
func WithLanguageModelStreamExtraFunc(fn LanguageModelStreamExtraFunc) LanguageModelOption {
	return func(l *languageModel) {
		l.streamExtraFunc = fn
	}
}

// WithLanguageModelUsageFunc sets the usage function for the language model.
func WithLanguageModelUsageFunc(fn LanguageModelUsageFunc) LanguageModelOption {
	return func(l *languageModel) {
		l.usageFunc = fn
	}
}

// WithLanguageModelStreamUsageFunc sets the stream usage function for the language model.
func WithLanguageModelStreamUsageFunc(fn LanguageModelStreamUsageFunc) LanguageModelOption {
	return func(l *languageModel) {
		l.streamUsageFunc = fn
	}
}

// WithLanguageModelToPromptFunc sets the to prompt function for the language model.
func WithLanguageModelToPromptFunc(fn LanguageModelToPromptFunc) LanguageModelOption {
	return func(l *languageModel) {
		l.toPromptFunc = fn
	}
}

// WithLanguageModelObjectMode sets the object generation mode.
func WithLanguageModelObjectMode(om fantasy.ObjectMode) LanguageModelOption {
	return func(l *languageModel) {
		// not supported
		if om == fantasy.ObjectModeJSON {
			om = fantasy.ObjectModeAuto
		}
		l.objectMode = om
	}
}

func newLanguageModel(modelID string, provider string, client openai.Client, opts ...LanguageModelOption) languageModel {
	model := languageModel{
		modelID:                    modelID,
		provider:                   provider,
		client:                     client,
		objectMode:                 fantasy.ObjectModeAuto,
		prepareCallFunc:            DefaultPrepareCallFunc,
		mapFinishReasonFunc:        DefaultMapFinishReasonFunc,
		usageFunc:                  DefaultUsageFunc,
		streamUsageFunc:            DefaultStreamUsageFunc,
		streamProviderMetadataFunc: DefaultStreamProviderMetadataFunc,
		toPromptFunc:               DefaultToPrompt,
	}

	for _, o := range opts {
		o(&model)
	}
	return model
}

type streamToolCall struct {
	id          string
	name        string
	arguments   string
	hasFinished bool
}

// Model implements fantasy.LanguageModel.
func (o languageModel) Model() string {
	return o.modelID
}

// Provider implements fantasy.LanguageModel.
func (o languageModel) Provider() string {
	return o.provider
}

func (o languageModel) prepareParams(call fantasy.Call) (*openai.ChatCompletionNewParams, []fantasy.CallWarning, error) {
	params := &openai.ChatCompletionNewParams{}
	messages, warnings := o.toPromptFunc(call.Prompt, o.provider, o.modelID)
	if call.TopK != nil {
		warnings = append(warnings, fantasy.CallWarning{
			Type:    fantasy.CallWarningTypeUnsupportedSetting,
			Setting: "top_k",
		})
	}

	if call.MaxOutputTokens != nil {
		params.MaxTokens = param.NewOpt(*call.MaxOutputTokens)
	}
	if call.Temperature != nil {
		params.Temperature = param.NewOpt(*call.Temperature)
	}
	if call.TopP != nil {
		params.TopP = param.NewOpt(*call.TopP)
	}
	if call.FrequencyPenalty != nil {
		params.FrequencyPenalty = param.NewOpt(*call.FrequencyPenalty)
	}
	if call.PresencePenalty != nil {
		params.PresencePenalty = param.NewOpt(*call.PresencePenalty)
	}

	if isReasoningModel(o.modelID) {
		// remove unsupported settings for reasoning models
		// see https://platform.openai.com/docs/guides/reasoning#limitations
		if call.Temperature != nil {
			params.Temperature = param.Opt[float64]{}
			warnings = append(warnings, fantasy.CallWarning{
				Type:    fantasy.CallWarningTypeUnsupportedSetting,
				Setting: "temperature",
				Details: "temperature is not supported for reasoning models",
			})
		}
		if call.TopP != nil {
			params.TopP = param.Opt[float64]{}
			warnings = append(warnings, fantasy.CallWarning{
				Type:    fantasy.CallWarningTypeUnsupportedSetting,
				Setting: "TopP",
				Details: "TopP is not supported for reasoning models",
			})
		}
		if call.FrequencyPenalty != nil {
			params.FrequencyPenalty = param.Opt[float64]{}
			warnings = append(warnings, fantasy.CallWarning{
				Type:    fantasy.CallWarningTypeUnsupportedSetting,
				Setting: "FrequencyPenalty",
				Details: "FrequencyPenalty is not supported for reasoning models",
			})
		}
		if call.PresencePenalty != nil {
			params.PresencePenalty = param.Opt[float64]{}
			warnings = append(warnings, fantasy.CallWarning{
				Type:    fantasy.CallWarningTypeUnsupportedSetting,
				Setting: "PresencePenalty",
				Details: "PresencePenalty is not supported for reasoning models",
			})
		}

		// reasoning models use max_completion_tokens instead of max_tokens
		if call.MaxOutputTokens != nil {
			if !params.MaxCompletionTokens.Valid() {
				params.MaxCompletionTokens = param.NewOpt(*call.MaxOutputTokens)
			}
			params.MaxTokens = param.Opt[int64]{}
		}
	}

	// Handle search preview models
	if isSearchPreviewModel(o.modelID) {
		if call.Temperature != nil {
			params.Temperature = param.Opt[float64]{}
			warnings = append(warnings, fantasy.CallWarning{
				Type:    fantasy.CallWarningTypeUnsupportedSetting,
				Setting: "temperature",
				Details: "temperature is not supported for the search preview models and has been removed.",
			})
		}
	}

	optionsWarnings, err := o.prepareCallFunc(o, params, call)
	if err != nil {
		return nil, nil, err
	}

	if len(optionsWarnings) > 0 {
		warnings = append(warnings, optionsWarnings...)
	}

	params.Messages = messages
	params.Model = o.modelID

	if len(call.Tools) > 0 {
		tools, toolChoice, toolWarnings := toOpenAiTools(call.Tools, call.ToolChoice)
		params.Tools = tools
		if toolChoice != nil {
			params.ToolChoice = *toolChoice
		}
		warnings = append(warnings, toolWarnings...)
	}
	return params, warnings, nil
}

// Generate implements fantasy.LanguageModel.
func (o languageModel) Generate(ctx context.Context, call fantasy.Call) (*fantasy.Response, error) {
	params, warnings, err := o.prepareParams(call)
	if err != nil {
		return nil, err
	}
	response, err := o.client.Chat.Completions.New(ctx, *params, callUARequestOptions(call)...)
	if err != nil {
		return nil, toProviderErr(err)
	}
	if response == nil {
		return nil, &fantasy.Error{Title: "no response", Message: "provider returned nil response"}
	}

	if len(response.Choices) == 0 {
		return nil, &fantasy.Error{Title: "no response", Message: "no response generated"}
	}
	choice := response.Choices[0]
	content := make([]fantasy.Content, 0, 1+len(choice.Message.ToolCalls)+len(choice.Message.Annotations))
	text := choice.Message.Content
	if text != "" {
		content = append(content, fantasy.TextContent{
			Text: text,
		})
	}
	if o.extraContentFunc != nil {
		extraContent := o.extraContentFunc(choice)
		content = append(content, extraContent...)
	}
	for _, tc := range choice.Message.ToolCalls {
		toolCallID := tc.ID
		content = append(content, fantasy.ToolCallContent{
			ProviderExecuted: false,
			ToolCallID:       toolCallID,
			ToolName:         tc.Function.Name,
			Input:            tc.Function.Arguments,
		})
	}
	for _, annotation := range choice.Message.Annotations {
		if annotation.Type == "url_citation" {
			content = append(content, fantasy.SourceContent{
				SourceType: fantasy.SourceTypeURL,
				ID:         uuid.NewString(),
				URL:        annotation.URLCitation.URL,
				Title:      annotation.URLCitation.Title,
			})
		}
	}

	usage, providerMetadata := o.usageFunc(*response)

	mappedFinishReason := o.mapFinishReasonFunc(choice.FinishReason)
	if len(choice.Message.ToolCalls) > 0 {
		mappedFinishReason = fantasy.FinishReasonToolCalls
	}
	return &fantasy.Response{
		Content:      content,
		Usage:        usage,
		FinishReason: mappedFinishReason,
		ProviderMetadata: fantasy.ProviderMetadata{
			Name: providerMetadata,
		},
		Warnings: warnings,
	}, nil
}

// Stream implements fantasy.LanguageModel.
func (o languageModel) Stream(ctx context.Context, call fantasy.Call) (fantasy.StreamResponse, error) {
	params, warnings, err := o.prepareParams(call)
	if err != nil {
		return nil, err
	}

	params.StreamOptions = openai.ChatCompletionStreamOptionsParam{
		IncludeUsage: openai.Bool(true),
	}

	stream := o.client.Chat.Completions.NewStreaming(ctx, *params, callUARequestOptions(call)...)
	isActiveText := false
	toolCalls := make(map[int64]streamToolCall)

	providerMetadata := fantasy.ProviderMetadata{
		Name: &ProviderMetadata{},
	}
	acc := openai.ChatCompletionAccumulator{}
	extraContext := make(map[string]any)
	var usage fantasy.Usage
	var finishReason string
	return func(yield func(fantasy.StreamPart) bool) {
		if len(warnings) > 0 {
			if !yield(fantasy.StreamPart{
				Type:     fantasy.StreamPartTypeWarnings,
				Warnings: warnings,
			}) {
				return
			}
		}
		for stream.Next() {
			chunk := stream.Current()
			acc.AddChunk(chunk)
			usage, providerMetadata = o.streamUsageFunc(chunk, extraContext, providerMetadata)
			if len(chunk.Choices) == 0 {
				continue
			}
			for _, choice := range chunk.Choices {
				if choice.FinishReason != "" {
					finishReason = choice.FinishReason
				}
				if choice.Delta.Content != "" {
					if !isActiveText {
						isActiveText = true
						if !yield(fantasy.StreamPart{
							Type: fantasy.StreamPartTypeTextStart,
							ID:   "0",
						}) {
							return
						}
					}
					if !yield(fantasy.StreamPart{
						Type:  fantasy.StreamPartTypeTextDelta,
						ID:    "0",
						Delta: choice.Delta.Content,
					}) {
						return
					}
				}
				if len(choice.Delta.ToolCalls) > 0 {
					if isActiveText {
						isActiveText = false
						if !yield(fantasy.StreamPart{
							Type: fantasy.StreamPartTypeTextEnd,
							ID:   "0",
						}) {
							return
						}
					}

					for _, toolCallDelta := range choice.Delta.ToolCalls {
						if existingToolCall, ok := toolCalls[toolCallDelta.Index]; ok {
							if existingToolCall.hasFinished {
								continue
							}
							if toolCallDelta.Function.Arguments != "" {
								existingToolCall.arguments += toolCallDelta.Function.Arguments
								if !yield(fantasy.StreamPart{
									Type:  fantasy.StreamPartTypeToolInputDelta,
									ID:    existingToolCall.id,
									Delta: toolCallDelta.Function.Arguments,
								}) {
									return
								}
							}
							toolCalls[toolCallDelta.Index] = existingToolCall
						} else {
							// Some provider like Ollama may send empty tool calls or miss some fields.
							// We'll skip when we don't have enough info and also assume sane defaults.
							if toolCallDelta.Function.Name == "" && toolCallDelta.Function.Arguments == "" {
								continue
							}
							toolCallDelta.Type = cmp.Or(toolCallDelta.Type, "function")
							toolCallDelta.ID = cmp.Or(toolCallDelta.ID, fmt.Sprintf("tool-call-%d", toolCallDelta.Index))

							if toolCallDelta.Type != "function" {
								yield(fantasy.StreamPart{
									Type:  fantasy.StreamPartTypeError,
									Error: &fantasy.Error{Title: "invalid provider response", Message: "expected 'function' type."},
								})
								return
							}

							// A new tool call index means any previously opened
							// calls are complete. Emit ToolInputEnd for each to
							// preserve per-call framing during streaming so that
							// consumers can reconstruct parallel tool calls from
							// the stream order without buffering. Close them in
							// index order for deterministic output.
							priorIndices := make([]int64, 0, len(toolCalls))
							for idx := range toolCalls {
								priorIndices = append(priorIndices, idx)
							}
							slices.Sort(priorIndices)
							for _, idx := range priorIndices {
								tc := toolCalls[idx]
								if idx >= toolCallDelta.Index || tc.hasFinished {
									continue
								}
								if tc.arguments == "" {
									tc.arguments = "{}"
									toolCalls[idx] = tc
								}
								if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeToolInputEnd, ID: tc.id}) {
									return
								}
								tc.hasFinished = true
								toolCalls[idx] = tc
							}

							if !yield(fantasy.StreamPart{
								Type:         fantasy.StreamPartTypeToolInputStart,
								ID:           toolCallDelta.ID,
								ToolCallName: toolCallDelta.Function.Name,
							}) {
								return
							}
							toolCalls[toolCallDelta.Index] = streamToolCall{
								id:        toolCallDelta.ID,
								name:      toolCallDelta.Function.Name,
								arguments: toolCallDelta.Function.Arguments,
							}

							if toolCallDelta.Function.Arguments != "" {
								if !yield(fantasy.StreamPart{
									Type:  fantasy.StreamPartTypeToolInputDelta,
									ID:    toolCallDelta.ID,
									Delta: toolCallDelta.Function.Arguments,
								}) {
									return
								}
							}
							continue
						}
					}
				}

				if o.streamExtraFunc != nil {
					updatedContext, shouldContinue := o.streamExtraFunc(chunk, yield, extraContext)
					if !shouldContinue {
						return
					}
					extraContext = updatedContext
				}
			}

			for _, choice := range chunk.Choices {
				if annotations := parseAnnotationsFromDelta(choice.Delta); len(annotations) > 0 {
					for _, annotation := range annotations {
						if annotation.Type == "url_citation" {
							if !yield(fantasy.StreamPart{
								Type:       fantasy.StreamPartTypeSource,
								ID:         uuid.NewString(),
								SourceType: fantasy.SourceTypeURL,
								URL:        annotation.URLCitation.URL,
								Title:      annotation.URLCitation.Title,
							}) {
								return
							}
						}
					}
				}
			}
		}
		err := stream.Err()
		if err == nil || errors.Is(err, io.EOF) {
			if isActiveText {
				isActiveText = false
				if !yield(fantasy.StreamPart{
					Type: fantasy.StreamPartTypeTextEnd,
					ID:   "0",
				}) {
					return
				}
			}

			// Finalize tool calls in index order after the stream completes.
			indices := make([]int64, 0, len(toolCalls))
			for idx := range toolCalls {
				indices = append(indices, idx)
			}
			slices.Sort(indices)
			for _, idx := range indices {
				tc := toolCalls[idx]
				if !tc.hasFinished {
					if tc.arguments == "" {
						tc.arguments = "{}"
						toolCalls[idx] = tc
					}
					if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeToolInputEnd, ID: tc.id}) {
						return
					}
				}
				if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeToolCall, ID: tc.id, ToolCallName: tc.name, ToolCallInput: tc.arguments}) {
					return
				}
				tc.hasFinished = true
				toolCalls[idx] = tc
			}

			if len(acc.Choices) > 0 {
				choice := acc.Choices[0]
				providerMetadata = o.streamProviderMetadataFunc(choice, providerMetadata)

				for _, annotation := range choice.Message.Annotations {
					if annotation.Type == "url_citation" {
						if !yield(fantasy.StreamPart{
							Type:       fantasy.StreamPartTypeSource,
							ID:         acc.ID,
							SourceType: fantasy.SourceTypeURL,
							URL:        annotation.URLCitation.URL,
							Title:      annotation.URLCitation.Title,
						}) {
							return
						}
					}
				}
			}
			mappedFinishReason := o.mapFinishReasonFunc(finishReason)
			if len(acc.Choices) > 0 {
				choice := acc.Choices[0]
				if len(choice.Message.ToolCalls) > 0 {
					mappedFinishReason = fantasy.FinishReasonToolCalls
				}
			}
			// Truncated stream: upstream closed without finish_reason and we
			// can't infer a tool-call turn. Surface as a retryable error so
			// the retry middleware re-runs the step.
			if finishReason == "" && mappedFinishReason != fantasy.FinishReasonToolCalls {
				err := ctx.Err()
				if err == nil {
					err = fantasy.NewIncompleteStreamError()
				}
				yield(fantasy.StreamPart{
					Type:  fantasy.StreamPartTypeError,
					Error: err,
				})
				return
			}
			yield(fantasy.StreamPart{
				Type:             fantasy.StreamPartTypeFinish,
				Usage:            usage,
				FinishReason:     mappedFinishReason,
				ProviderMetadata: providerMetadata,
			})
			return
		} else { //nolint: revive
			yield(fantasy.StreamPart{
				Type:  fantasy.StreamPartTypeError,
				Error: toProviderErr(err),
			})
			return
		}
	}, nil
}

func isReasoningModel(modelID string) bool {
	return strings.HasPrefix(modelID, "o1") || strings.Contains(modelID, "-o1") ||
		strings.HasPrefix(modelID, "o3") || strings.Contains(modelID, "-o3") ||
		strings.HasPrefix(modelID, "o4") || strings.Contains(modelID, "-o4") ||
		strings.HasPrefix(modelID, "oss") || strings.Contains(modelID, "-oss") ||
		strings.Contains(strings.ToLower(modelID), "gpt-5")
}

func isSearchPreviewModel(modelID string) bool {
	return strings.Contains(modelID, "search-preview")
}

func supportsFlexProcessing(modelID string) bool {
	return strings.HasPrefix(modelID, "o3") || strings.Contains(modelID, "-o3") ||
		strings.Contains(modelID, "o4-mini") ||
		strings.Contains(strings.ToLower(modelID), "gpt-5")
}

func supportsPriorityProcessing(modelID string) bool {
	return strings.Contains(strings.ToLower(modelID), "gpt-4") ||
		strings.Contains(strings.ToLower(modelID), "gpt-5") ||
		strings.HasPrefix(modelID, "o3") ||
		strings.Contains(modelID, "-o3") ||
		strings.Contains(modelID, "o4-mini")
}

func toOpenAiTools(tools []fantasy.Tool, toolChoice *fantasy.ToolChoice) (openAiTools []openai.ChatCompletionToolUnionParam, openAiToolChoice *openai.ChatCompletionToolChoiceOptionUnionParam, warnings []fantasy.CallWarning) {
	for _, tool := range tools {
		if tool.GetType() == fantasy.ToolTypeFunction {
			ft, ok := tool.(fantasy.FunctionTool)
			if !ok {
				continue
			}
			openAiTools = append(openAiTools, openai.ChatCompletionToolUnionParam{
				OfFunction: &openai.ChatCompletionFunctionToolParam{
					Function: shared.FunctionDefinitionParam{
						Name:        ft.Name,
						Description: param.NewOpt(ft.Description),
						Parameters:  openai.FunctionParameters(ft.InputSchema),
						Strict:      param.NewOpt(false),
					},
					Type: "function",
				},
			})
			continue
		}

		warnings = append(warnings, fantasy.CallWarning{
			Type:    fantasy.CallWarningTypeUnsupportedTool,
			Tool:    tool,
			Message: "tool is not supported",
		})
	}
	if toolChoice == nil {
		return openAiTools, openAiToolChoice, warnings
	}

	switch *toolChoice {
	case fantasy.ToolChoiceAuto:
		openAiToolChoice = &openai.ChatCompletionToolChoiceOptionUnionParam{
			OfAuto: param.NewOpt("auto"),
		}
	case fantasy.ToolChoiceNone:
		openAiToolChoice = &openai.ChatCompletionToolChoiceOptionUnionParam{
			OfAuto: param.NewOpt("none"),
		}
	case fantasy.ToolChoiceRequired:
		openAiToolChoice = &openai.ChatCompletionToolChoiceOptionUnionParam{
			OfAuto: param.NewOpt("required"),
		}
	default:
		openAiToolChoice = &openai.ChatCompletionToolChoiceOptionUnionParam{
			OfFunctionToolChoice: &openai.ChatCompletionNamedToolChoiceParam{
				Type: "function",
				Function: openai.ChatCompletionNamedToolChoiceFunctionParam{
					Name: string(*toolChoice),
				},
			},
		}
	}
	return openAiTools, openAiToolChoice, warnings
}

// parseAnnotationsFromDelta parses annotations from the raw JSON of a delta.
func parseAnnotationsFromDelta(delta openai.ChatCompletionChunkChoiceDelta) []openai.ChatCompletionMessageAnnotation {
	var annotations []openai.ChatCompletionMessageAnnotation

	// Parse the raw JSON to extract annotations
	var deltaData map[string]any
	if err := json.Unmarshal([]byte(delta.RawJSON()), &deltaData); err != nil {
		return annotations
	}

	// Check if annotations exist in the delta
	if annotationsData, ok := deltaData["annotations"].([]any); ok {
		for _, annotationData := range annotationsData {
			if annotationMap, ok := annotationData.(map[string]any); ok {
				if annotationType, ok := annotationMap["type"].(string); ok && annotationType == "url_citation" {
					if urlCitationData, ok := annotationMap["url_citation"].(map[string]any); ok {
						url, urlOk := urlCitationData["url"].(string)
						title, titleOk := urlCitationData["title"].(string)
						if urlOk && titleOk {
							annotation := openai.ChatCompletionMessageAnnotation{
								Type: "url_citation",
								URLCitation: openai.ChatCompletionMessageAnnotationURLCitation{
									URL:   url,
									Title: title,
								},
							}
							annotations = append(annotations, annotation)
						}
					}
				}
			}
		}
	}

	return annotations
}

// GenerateObject implements fantasy.LanguageModel.
func (o languageModel) GenerateObject(ctx context.Context, call fantasy.ObjectCall) (*fantasy.ObjectResponse, error) {
	switch o.objectMode {
	case fantasy.ObjectModeText:
		return object.GenerateWithText(ctx, o, call)
	case fantasy.ObjectModeTool:
		return object.GenerateWithTool(ctx, o, call)
	default:
		return o.generateObjectWithJSONMode(ctx, call)
	}
}

// StreamObject implements fantasy.LanguageModel.
func (o languageModel) StreamObject(ctx context.Context, call fantasy.ObjectCall) (fantasy.ObjectStreamResponse, error) {
	switch o.objectMode {
	case fantasy.ObjectModeTool:
		return object.StreamWithTool(ctx, o, call)
	case fantasy.ObjectModeText:
		return object.StreamWithText(ctx, o, call)
	default:
		return o.streamObjectWithJSONMode(ctx, call)
	}
}

func (o languageModel) generateObjectWithJSONMode(ctx context.Context, call fantasy.ObjectCall) (*fantasy.ObjectResponse, error) {
	jsonSchemaMap := schema.ToMap(call.Schema)

	addAdditionalPropertiesFalse(jsonSchemaMap)

	schemaName := call.SchemaName
	if schemaName == "" {
		schemaName = "response"
	}

	fantasyCall := fantasy.Call{
		Prompt:           call.Prompt,
		MaxOutputTokens:  call.MaxOutputTokens,
		Temperature:      call.Temperature,
		TopP:             call.TopP,
		PresencePenalty:  call.PresencePenalty,
		FrequencyPenalty: call.FrequencyPenalty,
		ProviderOptions:  call.ProviderOptions,
	}

	params, warnings, err := o.prepareParams(fantasyCall)
	if err != nil {
		return nil, err
	}

	params.ResponseFormat = openai.ChatCompletionNewParamsResponseFormatUnion{
		OfJSONSchema: &shared.ResponseFormatJSONSchemaParam{
			JSONSchema: shared.ResponseFormatJSONSchemaJSONSchemaParam{
				Name:        schemaName,
				Description: param.NewOpt(call.SchemaDescription),
				Schema:      jsonSchemaMap,
				Strict:      param.NewOpt(true),
			},
		},
	}

	response, err := o.client.Chat.Completions.New(ctx, *params, objectCallUARequestOptions(call)...)
	if err != nil {
		return nil, toProviderErr(err)
	}
	if len(response.Choices) == 0 {
		usage, _ := o.usageFunc(*response)
		return nil, &fantasy.NoObjectGeneratedError{
			RawText:      "",
			ParseError:   fmt.Errorf("no choices in response"),
			Usage:        usage,
			FinishReason: fantasy.FinishReasonUnknown,
		}
	}

	choice := response.Choices[0]
	jsonText := choice.Message.Content

	var obj any
	if call.RepairText != nil {
		obj, err = schema.ParseAndValidateWithRepair(ctx, jsonText, call.Schema, call.RepairText)
	} else {
		obj, err = schema.ParseAndValidate(jsonText, call.Schema)
	}

	usage, _ := o.usageFunc(*response)
	finishReason := o.mapFinishReasonFunc(choice.FinishReason)

	if err != nil {
		if nogErr, ok := err.(*fantasy.NoObjectGeneratedError); ok {
			nogErr.Usage = usage
			nogErr.FinishReason = finishReason
		}
		return nil, err
	}

	return &fantasy.ObjectResponse{
		Object:       obj,
		RawText:      jsonText,
		Usage:        usage,
		FinishReason: finishReason,
		Warnings:     warnings,
	}, nil
}

func (o languageModel) streamObjectWithJSONMode(ctx context.Context, call fantasy.ObjectCall) (fantasy.ObjectStreamResponse, error) {
	jsonSchemaMap := schema.ToMap(call.Schema)

	addAdditionalPropertiesFalse(jsonSchemaMap)

	schemaName := call.SchemaName
	if schemaName == "" {
		schemaName = "response"
	}

	fantasyCall := fantasy.Call{
		Prompt:           call.Prompt,
		MaxOutputTokens:  call.MaxOutputTokens,
		Temperature:      call.Temperature,
		TopP:             call.TopP,
		PresencePenalty:  call.PresencePenalty,
		FrequencyPenalty: call.FrequencyPenalty,
		ProviderOptions:  call.ProviderOptions,
	}

	params, warnings, err := o.prepareParams(fantasyCall)
	if err != nil {
		return nil, err
	}

	params.ResponseFormat = openai.ChatCompletionNewParamsResponseFormatUnion{
		OfJSONSchema: &shared.ResponseFormatJSONSchemaParam{
			JSONSchema: shared.ResponseFormatJSONSchemaJSONSchemaParam{
				Name:        schemaName,
				Description: param.NewOpt(call.SchemaDescription),
				Schema:      jsonSchemaMap,
				Strict:      param.NewOpt(true),
			},
		},
	}

	params.StreamOptions = openai.ChatCompletionStreamOptionsParam{
		IncludeUsage: openai.Bool(true),
	}

	stream := o.client.Chat.Completions.NewStreaming(ctx, *params, objectCallUARequestOptions(call)...)

	return func(yield func(fantasy.ObjectStreamPart) bool) {
		if len(warnings) > 0 {
			if !yield(fantasy.ObjectStreamPart{
				Type:     fantasy.ObjectStreamPartTypeObject,
				Warnings: warnings,
			}) {
				return
			}
		}

		var accumulated string
		var lastParsedObject any
		var usage fantasy.Usage
		var finishReason fantasy.FinishReason
		var sawFinishReason bool
		var providerMetadata fantasy.ProviderMetadata

		for stream.Next() {
			chunk := stream.Current()

			// Update usage
			usage, providerMetadata = o.streamUsageFunc(chunk, make(map[string]any), providerMetadata)

			if len(chunk.Choices) == 0 {
				continue
			}

			choice := chunk.Choices[0]
			if choice.FinishReason != "" {
				finishReason = o.mapFinishReasonFunc(choice.FinishReason)
				sawFinishReason = true
			}

			if choice.Delta.Content != "" {
				accumulated += choice.Delta.Content

				obj, state, parseErr := schema.ParsePartialJSON(accumulated)

				if state == schema.ParseStateSuccessful || state == schema.ParseStateRepaired {
					if err := schema.ValidateAgainstSchema(obj, call.Schema); err == nil {
						if !reflect.DeepEqual(obj, lastParsedObject) {
							if !yield(fantasy.ObjectStreamPart{
								Type:   fantasy.ObjectStreamPartTypeObject,
								Object: obj,
							}) {
								return
							}
							lastParsedObject = obj
						}
					}
				}

				if state == schema.ParseStateFailed && call.RepairText != nil {
					repairedText, repairErr := call.RepairText(ctx, accumulated, parseErr)
					if repairErr == nil {
						obj2, state2, _ := schema.ParsePartialJSON(repairedText)
						if (state2 == schema.ParseStateSuccessful || state2 == schema.ParseStateRepaired) &&
							schema.ValidateAgainstSchema(obj2, call.Schema) == nil {
							if !reflect.DeepEqual(obj2, lastParsedObject) {
								if !yield(fantasy.ObjectStreamPart{
									Type:   fantasy.ObjectStreamPartTypeObject,
									Object: obj2,
								}) {
									return
								}
								lastParsedObject = obj2
							}
						}
					}
				}
			}
		}

		err := stream.Err()
		if err != nil && !errors.Is(err, io.EOF) {
			yield(fantasy.ObjectStreamPart{
				Type:  fantasy.ObjectStreamPartTypeError,
				Error: toProviderErr(err),
			})
			return
		}

		if !sawFinishReason {
			err := ctx.Err()
			if err == nil {
				err = fantasy.NewIncompleteStreamError()
			}
			yield(fantasy.ObjectStreamPart{
				Type:  fantasy.ObjectStreamPartTypeError,
				Error: err,
			})
			return
		}

		if lastParsedObject != nil {
			yield(fantasy.ObjectStreamPart{
				Type:             fantasy.ObjectStreamPartTypeFinish,
				Usage:            usage,
				FinishReason:     finishReason,
				ProviderMetadata: providerMetadata,
			})
		} else {
			yield(fantasy.ObjectStreamPart{
				Type: fantasy.ObjectStreamPartTypeError,
				Error: &fantasy.NoObjectGeneratedError{
					RawText:      accumulated,
					ParseError:   fmt.Errorf("no valid object generated in stream"),
					Usage:        usage,
					FinishReason: finishReason,
				},
			})
		}
	}, nil
}

// addAdditionalPropertiesFalse recursively adds "additionalProperties": false to all object schemas.
// This is required by OpenAI's strict mode for structured outputs.
func addAdditionalPropertiesFalse(schema map[string]any) {
	if schema["type"] == "object" {
		if _, hasAdditional := schema["additionalProperties"]; !hasAdditional {
			schema["additionalProperties"] = false
		}

		// Recursively process nested properties
		if properties, ok := schema["properties"].(map[string]any); ok {
			for _, propValue := range properties {
				if propSchema, ok := propValue.(map[string]any); ok {
					addAdditionalPropertiesFalse(propSchema)
				}
			}
		}
	}

	// Handle array items
	if items, ok := schema["items"].(map[string]any); ok {
		addAdditionalPropertiesFalse(items)
	}
}
