package extproc

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

	core "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
	ext_proc "github.com/envoyproxy/go-control-plane/envoy/service/ext_proc/v3"
	typev3 "github.com/envoyproxy/go-control-plane/envoy/type/v3"
	"github.com/openai/openai-go"
	"github.com/tidwall/sjson"

	"github.com/vllm-project/semantic-router/src/semantic-router/pkg/observability/logging"
	"github.com/vllm-project/semantic-router/src/semantic-router/pkg/responseapi"
	"github.com/vllm-project/semantic-router/src/semantic-router/pkg/responsestore"
)

// ResponseAPIFilter handles Response API request/response translation.
// It translates /v1/responses requests to /v1/chat/completions format
// and translates responses back, managing conversation history via previous_response_id.
type ResponseAPIFilter struct {
	store      responsestore.ResponseStore
	translator *responseapi.Translator
	enabled    bool
}

// NewResponseAPIFilter creates a new Response API filter.
func NewResponseAPIFilter(store responsestore.ResponseStore) *ResponseAPIFilter {
	return &ResponseAPIFilter{
		store:      store,
		translator: responseapi.NewTranslator(),
		enabled:    store != nil && store.IsEnabled(),
	}
}

// IsEnabled returns whether Response API is enabled.
func (f *ResponseAPIFilter) IsEnabled() bool {
	return f.enabled
}

// ResponseAPIContext holds context for a Response API request during processing.
type ResponseAPIContext struct {
	// IsResponseAPIRequest indicates this is a /v1/responses request
	IsResponseAPIRequest bool

	// OriginalRequest is the parsed Response API request
	OriginalRequest *responseapi.ResponseAPIRequest

	// PreviousResponseID from the request (for conversation chaining)
	PreviousResponseID string

	// ConversationHistory fetched from store
	ConversationHistory []*responseapi.StoredResponse

	// GeneratedResponseID is the ID generated for this response
	GeneratedResponseID string

	// ConversationID is the determined ConversationID for this request.
	// Set early during TranslateRequest to ensure consistent tracking.
	// Sources (in priority order):
	//   1. Request's conversation_id field
	//   2. First response in conversation chain (via previous_response_id)
	//   3. Newly generated (for new conversations)
	ConversationID string

	// TranslatedBody is the Chat Completions request body after translation
	TranslatedBody []byte

	// HasImageGenerationTool is true when the request includes
	// tools: [{"type": "image_generation"}]. This signals that the client
	// explicitly requests image generation capability.
	HasImageGenerationTool bool

	// ImageGenToolParams holds the parameters from the image_generation tool
	// definition (model, quality, size, output_format, background, action).
	// Non-nil only when HasImageGenerationTool is true.
	ImageGenToolParams *responseapi.ImageGenerationToolParams
}

// detectImageGenTool scans the request tools for an image_generation entry and
// populates the context accordingly.
func detectImageGenTool(req *responseapi.ResponseAPIRequest, respCtx *ResponseAPIContext) {
	for i := range req.Tools {
		if req.Tools[i].Type == responseapi.ToolTypeImageGeneration {
			respCtx.HasImageGenerationTool = true
			respCtx.ImageGenToolParams = req.Tools[i].ExtractImageGenParams()
			if respCtx.ImageGenToolParams != nil {
				logging.Debugf("Response API: image_generation tool (model=%s, size=%s)", respCtx.ImageGenToolParams.Model, respCtx.ImageGenToolParams.Size)
			}
			return
		}
	}
}

// TranslateRequest translates a Response API request to Chat Completions format.
// Returns the translated request body and context, or nil if not a Response API request.
func (f *ResponseAPIFilter) TranslateRequest(ctx context.Context, body []byte) (*ResponseAPIContext, []byte, error) {
	if !f.enabled {
		return nil, nil, nil
	}

	var req responseapi.ResponseAPIRequest
	if err := json.Unmarshal(body, &req); err != nil {
		return nil, nil, nil
	}

	if len(req.Input) == 0 {
		return nil, nil, nil
	}

	respCtx := &ResponseAPIContext{
		IsResponseAPIRequest: true,
		OriginalRequest:      &req,
		PreviousResponseID:   req.PreviousResponseID,
		GeneratedResponseID:  responseapi.GenerateResponseID(),
	}

	detectImageGenTool(&req, respCtx)

	if req.PreviousResponseID != "" {
		history, err := f.store.GetConversationChain(ctx, req.PreviousResponseID)
		if err != nil && !errors.Is(err, responsestore.ErrNotFound) {
			logging.Warnf("Failed to fetch conversation history for %s: %v", req.PreviousResponseID, err)
		}
		respCtx.ConversationHistory = history
	}

	respCtx.ConversationID = f.determineConversationID(&req, respCtx.ConversationHistory)

	completionReq, err := f.translator.TranslateToCompletionRequest(&req, respCtx.ConversationHistory)
	if err != nil {
		return nil, nil, err
	}

	translatedBody, err := json.Marshal(completionReq)
	if err != nil {
		return nil, nil, err
	}

	// The SDK struct doesn't expose a Stream field (the SDK sets it via
	// request options internally). We inject it so the downstream pipeline
	// and the upstream backend see the correct "stream" flag.
	if req.Stream {
		if b, err := sjson.SetBytes(translatedBody, "stream", true); err == nil {
			translatedBody = b
		}
	}

	respCtx.TranslatedBody = translatedBody

	return respCtx, translatedBody, nil
}

// TranslateResponse translates a Chat Completions response back to Response API format.
func (f *ResponseAPIFilter) TranslateResponse(ctx context.Context, respCtx *ResponseAPIContext, body []byte) ([]byte, error) {
	if !f.enabled || respCtx == nil || !respCtx.IsResponseAPIRequest {
		return body, nil
	}

	if isResponseAPIErrorBody(body) {
		logging.Warnf("Response API: Backend returned error response, passing through")
		return body, nil
	}

	completionResp, ok := parseChatCompletionResponse(body)
	if !ok {
		return body, nil
	}

	responseAPIResp := f.buildTranslatedResponse(respCtx, completionResp)
	f.maybeStoreTranslatedResponse(ctx, respCtx, responseAPIResp)

	translatedBody, ok := marshalTranslatedResponse(body, responseAPIResp)
	if !ok {
		return body, nil
	}

	return translatedBody, nil
}

func isResponseAPIErrorBody(body []byte) bool {
	var errorCheck map[string]interface{}
	if err := json.Unmarshal(body, &errorCheck); err != nil {
		return false
	}
	_, hasError := errorCheck["error"]
	return hasError
}

func parseChatCompletionResponse(body []byte) (*openai.ChatCompletion, bool) {
	var completionResp openai.ChatCompletion
	if err := json.Unmarshal(body, &completionResp); err != nil {
		logging.Errorf("Response API: Failed to parse completion response: %v", err)
		return nil, false
	}

	if completionResp.ID == "" && len(completionResp.Choices) == 0 {
		logging.Warnf("Response API: Invalid completion response (no id or choices), passing through")
		return nil, false
	}

	return &completionResp, true
}

func (f *ResponseAPIFilter) buildTranslatedResponse(
	respCtx *ResponseAPIContext,
	completionResp *openai.ChatCompletion,
) *responseapi.ResponseAPIResponse {
	responseAPIResp := f.translator.TranslateToResponseAPIResponse(
		respCtx.OriginalRequest,
		completionResp,
		respCtx.PreviousResponseID,
	)

	responseAPIResp.ID = respCtx.GeneratedResponseID

	if respCtx.ConversationID != "" {
		responseAPIResp.ConversationID = respCtx.ConversationID
	}

	return responseAPIResp
}

func (f *ResponseAPIFilter) maybeStoreTranslatedResponse(
	ctx context.Context,
	respCtx *ResponseAPIContext,
	responseAPIResp *responseapi.ResponseAPIResponse,
) {
	shouldStore := respCtx.OriginalRequest.Store == nil || *respCtx.OriginalRequest.Store
	if shouldStore && f.store.IsEnabled() {
		stored := f.toStoredResponse(respCtx.OriginalRequest, responseAPIResp)
		if err := f.store.StoreResponse(ctx, stored); err != nil {
			logging.Warnf("Response API: Failed to store response: %v", err)
		}
	}
}

func marshalTranslatedResponse(
	originalBody []byte,
	responseAPIResp *responseapi.ResponseAPIResponse,
) ([]byte, bool) {
	translatedBody, err := json.Marshal(responseAPIResp)
	if err != nil {
		logging.Errorf("Response API: Failed to marshal response: %v", err)
		return originalBody, false
	}

	return translatedBody, true
}

// toStoredResponse converts request and response to a StoredResponse for storage.
func (f *ResponseAPIFilter) toStoredResponse(req *responseapi.ResponseAPIRequest, resp *responseapi.ResponseAPIResponse) *responseapi.StoredResponse {
	inputItems := parseResponseAPIInputItems(req.Input)

	return &responseapi.StoredResponse{
		ID:                 resp.ID,
		Object:             resp.Object,
		CreatedAt:          resp.CreatedAt,
		Model:              resp.Model,
		Status:             resp.Status,
		Input:              inputItems,
		Output:             resp.Output,
		OutputText:         resp.OutputText,
		PreviousResponseID: resp.PreviousResponseID,
		ConversationID:     resp.ConversationID,
		Usage:              resp.Usage,
		Instructions:       resp.Instructions,
		Metadata:           resp.Metadata,
	}
}

func parseResponseAPIInputItems(input json.RawMessage) []responseapi.InputItem {
	if len(input) == 0 {
		return nil
	}

	// Try parsing as array of input items first.
	var items []responseapi.InputItem
	if err := json.Unmarshal(input, &items); err == nil {
		for i := range items {
			if items[i].ID == "" {
				items[i].ID = responseapi.GenerateItemID()
			}
			if items[i].Status == "" {
				items[i].Status = responseapi.StatusCompleted
			}
		}
		return items
	}

	// Fallback: input can also be a string; store as a user message item.
	var inputStr string
	if err := json.Unmarshal(input, &inputStr); err == nil {
		return []responseapi.InputItem{{
			ID:      responseapi.GenerateItemID(),
			Type:    responseapi.ItemTypeMessage,
			Role:    responseapi.RoleUser,
			Content: input,
			Status:  responseapi.StatusCompleted,
		}}
	}

	return nil
}

// HandleGetResponse handles GET /v1/responses/{id} requests.
func (f *ResponseAPIFilter) HandleGetResponse(ctx context.Context, responseID string) (*ext_proc.ProcessingResponse, error) {
	if !f.enabled {
		return createResponseAPIError(404, "Response API not enabled"), nil
	}

	// Get response from store
	stored, err := f.store.GetResponse(ctx, responseID)
	if err != nil {
		if errors.Is(err, responsestore.ErrNotFound) {
			return createResponseAPIError(404, "Response not found: "+responseID), nil
		}
		logging.Errorf("Response API: Error getting response %s: %v", responseID, err)
		return createResponseAPIError(500, "Error retrieving response"), nil
	}

	// Convert to Response API format
	resp := f.storedToResponseAPIResponse(stored)

	// Marshal response
	body, err := json.Marshal(resp)
	if err != nil {
		logging.Errorf("Response API: Error marshaling response: %v", err)
		return createResponseAPIError(500, "Error serializing response"), nil
	}

	return createImmediateJSONResponse(200, body), nil
}

// HandleDeleteResponse handles DELETE /v1/responses/{id} requests.
func (f *ResponseAPIFilter) HandleDeleteResponse(ctx context.Context, responseID string) (*ext_proc.ProcessingResponse, error) {
	if !f.enabled {
		return createResponseAPIError(404, "Response API not enabled"), nil
	}

	// Delete response from store
	err := f.store.DeleteResponse(ctx, responseID)
	if err != nil {
		if errors.Is(err, responsestore.ErrNotFound) {
			return createResponseAPIError(404, "Response not found: "+responseID), nil
		}
		logging.Errorf("Response API: Error deleting response %s: %v", responseID, err)
		return createResponseAPIError(500, "Error deleting response"), nil
	}

	// Return deletion confirmation
	deleteResp := responseapi.DeleteResponseResult{
		ID:      responseID,
		Object:  "response.deleted",
		Deleted: true,
	}

	body, err := json.Marshal(deleteResp)
	if err != nil {
		logging.Errorf("Response API: Error marshaling delete response: %v", err)
		return createResponseAPIError(500, "Error serializing response"), nil
	}

	return createImmediateJSONResponse(200, body), nil
}

// HandleGetInputItems handles GET /v1/responses/{id}/input_items requests.
func (f *ResponseAPIFilter) HandleGetInputItems(ctx context.Context, responseID string) (*ext_proc.ProcessingResponse, error) {
	if !f.enabled {
		return createResponseAPIError(404, "Response API not enabled"), nil
	}

	// Get response from store
	stored, err := f.store.GetResponse(ctx, responseID)
	if err != nil {
		if errors.Is(err, responsestore.ErrNotFound) {
			return createResponseAPIError(404, "Response not found: "+responseID), nil
		}
		logging.Errorf("Response API: Error getting response %s: %v", responseID, err)
		return createResponseAPIError(500, "Error retrieving response"), nil
	}

	// Build input items list from stored response
	inputItems := f.buildInputItemsList(stored)

	// Create response with pagination structure
	listResp := responseapi.InputItemsListResponse{
		Object:  "list",
		Data:    inputItems,
		FirstID: "",
		LastID:  "",
		HasMore: false,
	}

	if len(inputItems) > 0 {
		listResp.FirstID = inputItems[0].ID
		listResp.LastID = inputItems[len(inputItems)-1].ID
	}

	body, err := json.Marshal(listResp)
	if err != nil {
		logging.Errorf("Response API: Error marshaling input items: %v", err)
		return createResponseAPIError(500, "Error serializing response"), nil
	}

	return createImmediateJSONResponse(200, body), nil
}

// buildInputItemsList builds the input items list from a stored response.
func (f *ResponseAPIFilter) buildInputItemsList(stored *responseapi.StoredResponse) []responseapi.InputItem {
	var items []responseapi.InputItem

	// Add instructions as system message if present
	if stored.Instructions != "" {
		contentParts := []responseapi.ContentPart{{Type: "input_text", Text: stored.Instructions}}
		contentJSON, _ := json.Marshal(contentParts)
		items = append(items, responseapi.InputItem{
			ID:      responseapi.GenerateItemID(),
			Type:    "message",
			Role:    "system",
			Content: contentJSON,
			Status:  "completed",
		})
	}

	// Add stored input items
	items = append(items, stored.Input...)

	return items
}

// storedToResponseAPIResponse converts a StoredResponse back to ResponseAPIResponse.
func (f *ResponseAPIFilter) storedToResponseAPIResponse(stored *responseapi.StoredResponse) *responseapi.ResponseAPIResponse {
	return &responseapi.ResponseAPIResponse{
		ID:                 stored.ID,
		Object:             "response",
		CreatedAt:          stored.CreatedAt,
		Model:              stored.Model,
		Status:             stored.Status,
		Output:             stored.Output,
		OutputText:         stored.OutputText,
		PreviousResponseID: stored.PreviousResponseID,
		ConversationID:     stored.ConversationID,
		Usage:              stored.Usage,
		Instructions:       stored.Instructions,
		Metadata:           stored.Metadata,
	}
}

// determineConversationID determines the ConversationID for this request.
// Priority order:
//  1. Request's conversation_id field (explicit)
//  2. First response in conversation chain (continuation via previous_response_id)
//  3. Newly generated (new conversation)
func (f *ResponseAPIFilter) determineConversationID(req *responseapi.ResponseAPIRequest, history []*responseapi.StoredResponse) string {
	// Priority 1: Request explicitly provides conversation_id
	if req.ConversationID != "" {
		return req.ConversationID
	}

	// Priority 2: Get from conversation history (continuation)
	if len(history) > 0 {
		// Find ConversationID from the first response in the chain
		firstResponse := history[0]
		if firstResponse.ConversationID != "" {
			return firstResponse.ConversationID
		}
	}

	// Priority 3: Generate new ConversationID (new conversation)
	return responseapi.GenerateConversationID()
}

// createResponseAPIError creates an error response in OpenAI format.
func createResponseAPIError(statusCode int, message string) *ext_proc.ProcessingResponse {
	errorResp := map[string]interface{}{
		"error": map[string]interface{}{
			"message": message,
			"type":    "invalid_request_error",
			"code":    statusCode,
		},
	}

	body, _ := json.Marshal(errorResp)
	return createImmediateJSONResponse(statusCode, body)
}

// createImmediateJSONResponse creates an immediate response with JSON body.
func createImmediateJSONResponse(statusCode int, body []byte) *ext_proc.ProcessingResponse {
	return &ext_proc.ProcessingResponse{
		Response: &ext_proc.ProcessingResponse_ImmediateResponse{
			ImmediateResponse: &ext_proc.ImmediateResponse{
				Status: &typev3.HttpStatus{
					Code: statusCodeToEnumForResponseAPI(statusCode),
				},
				Headers: &ext_proc.HeaderMutation{
					SetHeaders: []*core.HeaderValueOption{
						{
							Header: &core.HeaderValue{
								Key:      "content-type",
								RawValue: []byte("application/json"),
							},
						},
					},
				},
				Body: body,
			},
		},
	}
}

// statusCodeToEnumForResponseAPI converts HTTP status code to Envoy enum.
func statusCodeToEnumForResponseAPI(statusCode int) typev3.StatusCode {
	switch statusCode {
	case 200:
		return typev3.StatusCode_OK
	case 400:
		return typev3.StatusCode_BadRequest
	case 404:
		return typev3.StatusCode_NotFound
	case 500:
		return typev3.StatusCode_InternalServerError
	default:
		return typev3.StatusCode_OK
	}
}
