package openai

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"regexp"
	"strings"

	"go.uber.org/zap"
)

var geminiPathToken = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]*$`)
var geminiCallerModelLabel = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9 ._-]{0,127}$`)
var geminiToolCallID = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$`)

const geminiRejectionLogMessage = "edge_gemini_rejection"

type geminiRejectionClass string

const (
	geminiRejectionPreIngress   geminiRejectionClass = "pre_ingress"
	geminiRejectionProviderHTTP geminiRejectionClass = "provider_http"
)

func isGeminiRequest(r *http.Request) bool {
	return r != nil && strings.HasPrefix(r.URL.Path, geminiPathPrefix)
}

func writeGeminiError(w http.ResponseWriter, status int, code, message string) {
	writeJSON(w, status, geminiErrorResponse{Error: geminiErrorBody{
		Code: status, Message: message, Status: code,
	}})
}

func (s *Server) observeGeminiRejection(class geminiRejectionClass, status int) {
	s.logger.Info(
		geminiRejectionLogMessage,
		zap.String("surface", "gemini"),
		zap.String("bridge", "chat"),
		zap.String("rejection_class", string(class)),
		zap.Int("http_status", status),
	)
}

func (s *Server) writeGeminiPreIngressError(w http.ResponseWriter, status int, code, message string) {
	s.observeGeminiRejection(geminiRejectionPreIngress, status)
	writeGeminiError(w, status, code, message)
}

func (s *Server) handleGeminiStreamGenerateContent(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		s.writeGeminiPreIngressError(w, http.StatusMethodNotAllowed, "INVALID_ARGUMENT", "method not allowed")
		return
	}
	routeID, callerModel, err := parseGeminiStreamPath(r)
	if err != nil {
		s.writeGeminiPreIngressError(w, http.StatusBadRequest, "INVALID_ARGUMENT", "request path is invalid")
		return
	}
	defer r.Body.Close()
	body, err := readOpenAIIngressBody(w, r, s.maxIngressSnapshotBytes())
	if err != nil {
		s.writeGeminiPreIngressError(w, http.StatusBadRequest, "INVALID_ARGUMENT", "request body is invalid")
		return
	}
	chatBody, err := prepareGeminiChatBridge(body, routeID)
	if err != nil {
		s.writeGeminiPreIngressError(w, http.StatusBadRequest, "INVALID_ARGUMENT", "request body is invalid")
		return
	}

	internal := r.Clone(r.Context())
	internal.URL.Path = "/v1/chat/completions"
	internal.URL.RawPath = ""
	internal.URL.RawQuery = ""
	internal.RequestURI = "/v1/chat/completions"
	internal.Body = io.NopCloser(bytes.NewReader(chatBody))
	internal.ContentLength = int64(len(chatBody))
	internal.Header = r.Header.Clone()
	internal.Header.Del("Authorization")
	internal.Header.Del("X-Goog-Api-Key")
	internal.Header.Set("Content-Type", "application/json")

	bridge := newGeminiBridgeResponseWriter(w, callerModel, func(status int) {
		s.observeGeminiRejection(geminiRejectionProviderHTTP, status)
	})
	s.handleChatCompletions(bridge, internal)
	bridge.Finish()
}

func parseGeminiStreamPath(r *http.Request) (string, string, error) {
	if r == nil || r.URL == nil {
		return "", "", fmt.Errorf("missing URL")
	}
	parts := strings.Split(strings.TrimPrefix(r.URL.Path, geminiPathPrefix), "/")
	if len(parts) != 4 || parts[1] != "v1beta" || parts[2] != "models" {
		return "", "", fmt.Errorf("unexpected path")
	}
	const suffix = ":streamGenerateContent"
	if !strings.HasSuffix(parts[3], suffix) {
		return "", "", fmt.Errorf("unexpected method")
	}
	routeID := parts[0]
	callerModel := strings.TrimSuffix(parts[3], suffix)
	if !geminiPathToken.MatchString(routeID) || !geminiCallerModelLabel.MatchString(callerModel) || strings.TrimSpace(callerModel) != callerModel {
		return "", "", fmt.Errorf("invalid path token")
	}
	query := r.URL.Query()
	if len(query) != 1 || len(query["alt"]) != 1 || query.Get("alt") != "sse" {
		return "", "", fmt.Errorf("alt=sse is required")
	}
	return routeID, callerModel, nil
}

func prepareGeminiChatBridge(body []byte, routeID string) ([]byte, error) {
	if err := validateJSONMembers(body); err != nil {
		return nil, err
	}
	decoder := json.NewDecoder(bytes.NewReader(body))
	decoder.DisallowUnknownFields()
	var req geminiRequest
	if err := decoder.Decode(&req); err != nil {
		return nil, err
	}
	if err := requireGeminiJSONEOF(decoder); err != nil {
		return nil, err
	}
	if len(req.Contents) == 0 {
		return nil, fmt.Errorf("contents are required")
	}
	chat := map[string]any{
		"model": routeID, "stream": true,
		"stream_options": map[string]any{"include_usage": true},
	}
	messages := make([]map[string]any, 0, len(req.Contents)+1)
	if req.SystemInstruction != nil {
		text, err := geminiTextOnly(*req.SystemInstruction)
		if err != nil || strings.TrimSpace(text) == "" {
			return nil, fmt.Errorf("systemInstruction is invalid")
		}
		messages = append(messages, map[string]any{"role": "system", "content": text})
	}
	pendingCalls := newGeminiPendingCalls()
	for contentIndex, content := range req.Contents {
		converted, err := geminiContentToChat(content, contentIndex, pendingCalls)
		if err != nil {
			return nil, err
		}
		messages = append(messages, converted...)
	}
	chat["messages"] = messages
	if config := req.GenerationConfig; config != nil {
		if config.CandidateCount != nil && *config.CandidateCount != 1 {
			return nil, fmt.Errorf("candidateCount must be one")
		}
		if config.MaxOutputTokens != nil {
			if *config.MaxOutputTokens <= 0 {
				return nil, fmt.Errorf("maxOutputTokens must be positive")
			}
			chat["max_tokens"] = *config.MaxOutputTokens
		}
		if len(config.StopSequences) > 0 {
			chat["stop"] = config.StopSequences
		}
		if config.Temperature != nil {
			if *config.Temperature < 0 || *config.Temperature > 2 {
				return nil, fmt.Errorf("temperature is invalid")
			}
		}
		if config.TopK != nil {
			if *config.TopK <= 0 {
				return nil, fmt.Errorf("topK is invalid")
			}
		}
		if config.TopP != nil {
			if *config.TopP < 0 || *config.TopP > 1 {
				return nil, fmt.Errorf("topP is invalid")
			}
		}
		if thinking := config.ThinkingConfig; thinking != nil {
			googleThinking := make(map[string]any)
			if thinking.IncludeThoughts != nil {
				googleThinking["include_thoughts"] = *thinking.IncludeThoughts
			}
			if thinking.ThinkingBudget != nil {
				if *thinking.ThinkingBudget < -1 {
					return nil, fmt.Errorf("thinkingBudget is invalid")
				}
				googleThinking["thinking_budget"] = *thinking.ThinkingBudget
			}
			if len(googleThinking) > 0 {
				// Gemini's OpenAI-compatible endpoint accepts native options only
				// below extra_body.google; generic think/include_reasoning fields
				// are rejected by that endpoint.
				chat["extra_body"] = map[string]any{
					"google": map[string]any{"thinking_config": googleThinking},
				}
			}
		}
		responseSchema, err := geminiExclusiveSchema(config.ResponseSchema, config.ResponseJSONSchema)
		if err != nil {
			return nil, fmt.Errorf("response schema is invalid")
		}
		switch config.ResponseMIMEType {
		case "":
			if responseSchema != nil {
				return nil, fmt.Errorf("response MIME type is required")
			}
		case "text/plain":
			if responseSchema != nil {
				return nil, fmt.Errorf("text response schema is invalid")
			}
		case "application/json":
			if responseSchema == nil {
				chat["response_format"] = map[string]any{"type": "json_object"}
			} else {
				chat["response_format"] = map[string]any{
					"type": "json_schema",
					"json_schema": map[string]any{
						"name": "agy_response", "strict": true, "schema": responseSchema,
					},
				}
			}
		default:
			return nil, fmt.Errorf("response MIME type is invalid")
		}
	}
	if len(req.Tools) > 0 {
		tools := make([]map[string]any, 0)
		for _, group := range req.Tools {
			if len(group.FunctionDeclarations) == 0 {
				return nil, fmt.Errorf("functionDeclarations are required")
			}
			for _, declaration := range group.FunctionDeclarations {
				if !geminiPathToken.MatchString(declaration.Name) {
					return nil, fmt.Errorf("function declaration is invalid")
				}
				schema, err := geminiExclusiveSchema(
					declaration.Parameters, declaration.ParametersJSONSchema,
				)
				if err != nil || schema == nil {
					return nil, fmt.Errorf("function schema is invalid")
				}
				if _, err := geminiExclusiveSchema(
					declaration.Response, declaration.ResponseJSONSchema,
				); err != nil {
					return nil, fmt.Errorf("function response schema is invalid")
				}
				function := map[string]any{"name": declaration.Name, "parameters": schema}
				if declaration.Description != "" {
					function["description"] = declaration.Description
				}
				tools = append(tools, map[string]any{"type": "function", "function": function})
			}
		}
		chat["tools"] = tools
	}
	if req.ToolConfig != nil && req.ToolConfig.FunctionCallingConfig != nil {
		switch strings.ToUpper(req.ToolConfig.FunctionCallingConfig.Mode) {
		case "", "AUTO":
			chat["tool_choice"] = "auto"
		case "ANY":
			chat["tool_choice"] = "required"
		case "NONE":
			chat["tool_choice"] = "none"
		default:
			return nil, fmt.Errorf("function calling mode is invalid")
		}
	}
	return json.Marshal(chat)
}

func geminiExclusiveSchema(first, second json.RawMessage) (map[string]any, error) {
	if len(first) > 0 && len(second) > 0 {
		return nil, fmt.Errorf("schema alternatives conflict")
	}
	raw := first
	if len(raw) == 0 {
		raw = second
	}
	if len(raw) == 0 {
		return nil, nil
	}
	var schema map[string]any
	if err := json.Unmarshal(raw, &schema); err != nil || schema == nil {
		return nil, fmt.Errorf("schema is not an object")
	}
	if err := normalizeGeminiSchemaTypes(schema); err != nil {
		return nil, err
	}
	return schema, nil
}

func normalizeGeminiSchemaTypes(value any) error {
	switch typed := value.(type) {
	case map[string]any:
		for key, item := range typed {
			if key == "type" {
				name, ok := item.(string)
				if !ok {
					return fmt.Errorf("schema type is invalid")
				}
				normalized := strings.ToLower(name)
				switch normalized {
				case "null", "boolean", "object", "array", "number", "integer", "string":
					typed[key] = normalized
				default:
					return fmt.Errorf("schema type is invalid")
				}
				continue
			}
			if err := normalizeGeminiSchemaTypes(item); err != nil {
				return err
			}
		}
	case []any:
		for _, item := range typed {
			if err := normalizeGeminiSchemaTypes(item); err != nil {
				return err
			}
		}
	}
	return nil
}

func geminiTextOnly(content geminiContent) (string, error) {
	// Gemini represents systemInstruction as Content and official agy 1.1.12
	// labels that Content with the API-native "user" role.
	if content.Role != "" && content.Role != "user" {
		return "", fmt.Errorf("role is invalid")
	}
	texts := make([]string, 0, len(content.Parts))
	for _, part := range content.Parts {
		if part.Text == nil || part.FunctionCall != nil || part.FunctionResponse != nil || part.Thought || part.ThoughtSignature != "" {
			return "", fmt.Errorf("only text is supported")
		}
		texts = append(texts, *part.Text)
	}
	return strings.Join(texts, "\n"), nil
}

type geminiPendingCall struct {
	id       string
	name     string
	explicit bool
}

type geminiPendingCalls struct {
	byID   map[string]geminiPendingCall
	byName map[string][]string
}

func newGeminiPendingCalls() *geminiPendingCalls {
	return &geminiPendingCalls{byID: make(map[string]geminiPendingCall), byName: make(map[string][]string)}
}

func (p *geminiPendingCalls) add(name, id string, explicit bool) error {
	if _, duplicate := p.byID[id]; duplicate {
		return fmt.Errorf("duplicate functionCall id")
	}
	p.byID[id] = geminiPendingCall{id: id, name: name, explicit: explicit}
	p.byName[name] = append(p.byName[name], id)
	return nil
}

func (p *geminiPendingCalls) consume(name, id string) (string, error) {
	if id != "" {
		if !geminiToolCallID.MatchString(id) {
			return "", fmt.Errorf("functionResponse id is invalid")
		}
		call, ok := p.byID[id]
		if !ok || call.name != name {
			return "", fmt.Errorf("functionResponse has no matching call")
		}
		p.remove(call)
		return id, nil
	}
	ids := p.byName[name]
	if len(ids) == 0 {
		return "", fmt.Errorf("functionResponse has no matching call")
	}
	call, ok := p.byID[ids[0]]
	if !ok || call.explicit {
		return "", fmt.Errorf("functionResponse id is required")
	}
	p.remove(call)
	return call.id, nil
}

func (p *geminiPendingCalls) remove(call geminiPendingCall) {
	delete(p.byID, call.id)
	ids := p.byName[call.name]
	for index, id := range ids {
		if id != call.id {
			continue
		}
		ids = append(ids[:index], ids[index+1:]...)
		break
	}
	if len(ids) == 0 {
		delete(p.byName, call.name)
	} else {
		p.byName[call.name] = ids
	}
}

func geminiContentToChat(content geminiContent, contentIndex int, pending *geminiPendingCalls) ([]map[string]any, error) {
	role := strings.ToLower(strings.TrimSpace(content.Role))
	if role != "user" && role != "model" {
		return nil, fmt.Errorf("content role is invalid")
	}
	if len(content.Parts) == 0 {
		return nil, fmt.Errorf("content parts are required")
	}
	var texts []string
	var reasoning []string
	var toolCalls []any
	var toolMessages []map[string]any
	for partIndex, part := range content.Parts {
		set := 0
		if part.Text != nil {
			set++
		}
		if part.FunctionCall != nil {
			set++
		}
		if part.FunctionResponse != nil {
			set++
		}
		if set != 1 {
			return nil, fmt.Errorf("content part is invalid")
		}
		if part.Text != nil {
			if role == "model" && part.Thought {
				reasoning = append(reasoning, *part.Text)
			} else if part.Thought {
				return nil, fmt.Errorf("user thought is invalid")
			} else {
				texts = append(texts, *part.Text)
			}
			continue
		}
		if part.FunctionCall != nil {
			if role != "model" || !geminiPathToken.MatchString(part.FunctionCall.Name) {
				return nil, fmt.Errorf("functionCall is invalid")
			}
			var args map[string]any
			if json.Unmarshal(part.FunctionCall.Args, &args) != nil {
				return nil, fmt.Errorf("functionCall args are invalid")
			}
			callID := part.FunctionCall.ID
			explicitID := callID != ""
			if !explicitID {
				callID = fmt.Sprintf("gemini_call_%d_%d", contentIndex, partIndex)
			} else if !geminiToolCallID.MatchString(callID) {
				return nil, fmt.Errorf("functionCall id is invalid")
			}
			if err := pending.add(part.FunctionCall.Name, callID, explicitID); err != nil {
				return nil, err
			}
			call := map[string]any{"id": callID, "type": "function", "function": map[string]any{"name": part.FunctionCall.Name, "arguments": string(part.FunctionCall.Args)}}
			if part.ThoughtSignature != "" {
				call["extra_content"] = openAIChatThoughtSignature(part.ThoughtSignature)
			}
			toolCalls = append(toolCalls, call)
			continue
		}
		response := part.FunctionResponse
		if !geminiPathToken.MatchString(response.Name) {
			return nil, fmt.Errorf("functionResponse is invalid")
		}
		if !json.Valid(response.Response) {
			return nil, fmt.Errorf("functionResponse has no matching call")
		}
		callID, err := pending.consume(response.Name, response.ID)
		if err != nil {
			return nil, err
		}
		toolMessages = append(toolMessages, map[string]any{"role": "tool", "tool_call_id": callID, "content": string(response.Response)})
	}
	if role == "model" {
		if len(toolMessages) > 0 {
			if len(texts) > 0 || len(reasoning) > 0 || len(toolCalls) > 0 {
				return nil, fmt.Errorf("model functionResponse cannot be mixed with assistant content")
			}
			return toolMessages, nil
		}
		message := map[string]any{"role": "assistant", "content": strings.Join(texts, "\n")}
		if len(reasoning) > 0 {
			message["reasoning_content"] = strings.Join(reasoning, "")
		}
		if len(toolCalls) > 0 {
			message["tool_calls"] = toolCalls
		}
		return []map[string]any{message}, nil
	}
	if len(texts) > 0 {
		toolMessages = append(toolMessages, map[string]any{"role": "user", "content": strings.Join(texts, "\n")})
	}
	if len(toolMessages) == 0 {
		return nil, fmt.Errorf("user content is empty")
	}
	return toolMessages, nil
}

func validateJSONMembers(raw []byte) error {
	dec := json.NewDecoder(bytes.NewReader(raw))
	if err := validateGeminiJSONValue(dec); err != nil {
		return err
	}
	return requireGeminiJSONEOF(dec)
}

func validateGeminiJSONValue(dec *json.Decoder) error {
	token, err := dec.Token()
	if err != nil {
		return err
	}
	delim, ok := token.(json.Delim)
	if !ok {
		return nil
	}
	switch delim {
	case '{':
		seen := map[string]struct{}{}
		for dec.More() {
			keyToken, err := dec.Token()
			if err != nil {
				return err
			}
			key, ok := keyToken.(string)
			if !ok {
				return fmt.Errorf("object key is invalid")
			}
			if _, duplicate := seen[key]; duplicate {
				return fmt.Errorf("duplicate member")
			}
			seen[key] = struct{}{}
			if err := validateGeminiJSONValue(dec); err != nil {
				return err
			}
		}
		closeToken, err := dec.Token()
		if err != nil || closeToken != json.Delim('}') {
			return fmt.Errorf("object is invalid")
		}
	case '[':
		for dec.More() {
			if err := validateGeminiJSONValue(dec); err != nil {
				return err
			}
		}
		closeToken, err := dec.Token()
		if err != nil || closeToken != json.Delim(']') {
			return fmt.Errorf("array is invalid")
		}
	default:
		return fmt.Errorf("JSON delimiter is invalid")
	}
	return nil
}

func requireGeminiJSONEOF(dec *json.Decoder) error {
	var extra any
	if err := dec.Decode(&extra); err != io.EOF {
		return fmt.Errorf("trailing JSON value")
	}
	return nil
}