package gateway

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"time"

	"github.com/rose_cat707/luminary/internal/auth"
	"github.com/rose_cat707/luminary/internal/limiter"
	"github.com/rose_cat707/luminary/internal/model"
	"github.com/rose_cat707/luminary/internal/pricing"
	"github.com/rose_cat707/luminary/internal/provider"
	"github.com/rose_cat707/luminary/internal/provider/anthropic"
	"github.com/rose_cat707/luminary/internal/provider/comfyui"
	"github.com/rose_cat707/luminary/internal/provider/openai"
	"github.com/rose_cat707/luminary/internal/settings"
	"github.com/rose_cat707/luminary/internal/store/cache"
)

type Service struct {
	cache        *cache.Store
	settings     *settings.Runtime
	openai       *openai.Client
	anthropic    *anthropic.Client
	streamClient *provider.HTTPClientWrapper
}

func New(c *cache.Store, rt *settings.Runtime) *Service {
	return &Service{
		cache:        c,
		settings:     rt,
		openai:       openai.New(),
		anthropic:    anthropic.New(),
		streamClient: &provider.HTTPClientWrapper{Client: provider.NewStreamHTTPClient()},
	}
}

func (s *Service) ResolveIngressKey(plainKey string) (*model.IngressKey, error) {
	k, ok := s.cache.LookupIngressKey(plainKey)
	if !ok {
		return nil, fmt.Errorf("invalid api key")
	}
	if k.ExpiresAt != nil && k.ExpiresAt.Before(time.Now()) {
		return nil, fmt.Errorf("api key expired")
	}
	return k, nil
}

func (s *Service) providerConfig(p *model.Provider) provider.ProviderConfig {
	return provider.ProviderConfig{
		Type:          p.Type,
		Category:      p.Category,
		BaseURL:       p.BaseURL,
		EndpointPaths: provider.ParseEndpointPaths(p.EndpointPathsJSON),
	}
}

func (s *Service) DecryptAPIKey(enc string) (string, error) {
	return auth.Decrypt(enc, s.settings.EncryptionKey())
}

func (s *Service) providerAllowsEndpoint(p *model.Provider, endpoint provider.EndpointType) bool {
	allowed := provider.ParseAllowedEndpoints(p.AllowedEndpointsJSON, p.Category)
	return provider.IsEndpointAllowed(allowed, endpoint)
}

func (s *Service) canServeEndpoint(p *model.Provider, endpoint provider.EndpointType) bool {
	if s.providerAllowsEndpoint(p, endpoint) {
		return true
	}
	return endpoint == provider.EndpointResponses && s.providerAllowsEndpoint(p, provider.EndpointChatCompletion)
}

func (s *Service) shouldAdaptResponses(p *model.Provider) bool {
	return !s.providerAllowsEndpoint(p, provider.EndpointResponses) && s.providerAllowsEndpoint(p, provider.EndpointChatCompletion)
}

func (s *Service) forward(ctx context.Context, p *model.Provider, apiKey string, endpoint provider.EndpointType, body []byte) (*provider.ChatResponse, error) {
	var err error
	body, err = s.rewriteBodyForProvider(p.ID, body)
	if err != nil {
		return nil, err
	}
	actualEndpoint := endpoint
	actualBody := body
	adaptResponses := false

	if endpoint == provider.EndpointResponses {
		if s.shouldAdaptResponses(p) {
			adapted, err := provider.ResponsesToChatCompletion(body)
			if err != nil {
				return nil, err
			}
			actualBody = adapted
			actualEndpoint = provider.EndpointChatCompletion
			adaptResponses = true
		} else if !s.providerAllowsEndpoint(p, endpoint) {
			return nil, fmt.Errorf("endpoint %s not enabled on provider %s", endpoint, p.Name)
		}
	} else if !s.providerAllowsEndpoint(p, endpoint) {
		return nil, fmt.Errorf("endpoint %s not enabled on provider %s", endpoint, p.Name)
	}

	resp, err := s.forwardRaw(ctx, p, apiKey, actualEndpoint, actualBody)
	if err != nil {
		return nil, err
	}

	if endpoint == provider.EndpointResponses && !adaptResponses && resp != nil && resp.StatusCode == http.StatusNotFound && s.providerAllowsEndpoint(p, provider.EndpointChatCompletion) {
		adaptedBody, adaptErr := provider.ResponsesToChatCompletion(body)
		if adaptErr == nil {
			if resp2, err2 := s.forwardRaw(ctx, p, apiKey, provider.EndpointChatCompletion, adaptedBody); err2 == nil && resp2 != nil {
				resp = resp2
				adaptResponses = true
			}
		}
	}

	if adaptResponses && resp != nil && resp.StatusCode < 400 {
		converted, convErr := provider.ChatCompletionToResponses(resp.Body)
		if convErr == nil {
			resp.Body = converted
		}
	}
	return resp, nil
}

func (s *Service) forwardRaw(ctx context.Context, p *model.Provider, apiKey string, endpoint provider.EndpointType, body []byte) (*provider.ChatResponse, error) {
	cfg := s.providerConfig(p)
	if p.Type == model.ProviderAnthropic && endpoint == provider.EndpointChatCompletion {
		return s.anthropic.ForwardWithConfig(ctx, apiKey, cfg, endpoint, body)
	}
	return s.openai.ForwardWithConfig(ctx, apiKey, cfg, endpoint, body)
}

func (s *Service) forwardStream(ctx context.Context, p *model.Provider, apiKey string, endpoint provider.EndpointType, body []byte, w io.Writer, flusher http.Flusher) (*provider.StreamResult, error) {
	var err error
	body, err = s.rewriteBodyForProvider(p.ID, body)
	if err != nil {
		return nil, err
	}
	actualEndpoint := endpoint
	actualBody := body
	adaptResponses := false

	if endpoint == provider.EndpointResponses {
		if s.shouldAdaptResponses(p) {
			adapted, err := provider.ResponsesToChatCompletion(body)
			if err != nil {
				return nil, err
			}
			actualBody = adapted
			actualEndpoint = provider.EndpointChatCompletion
			adaptResponses = true
		} else if !s.providerAllowsEndpoint(p, endpoint) {
			return nil, fmt.Errorf("endpoint %s not enabled", endpoint)
		}
	} else if !s.providerAllowsEndpoint(p, endpoint) {
		return nil, fmt.Errorf("endpoint %s not enabled", endpoint)
	}

	cfg := s.providerConfig(p)
	url := cfg.ResolveEndpointURL(actualEndpoint)

	if adaptResponses {
		pr, pw := io.Pipe()
		errCh := make(chan error, 1)
		go func() {
			_, err := provider.ForwardStreamHTTP(ctx, s.streamClient.Client, "POST", url, apiKey, p.Type, actualBody, pw, nil)
			_ = pw.Close()
			errCh <- err
		}()
		result, err := provider.AdaptResponsesStream(body, pr, w, flusher)
		if streamErr := <-errCh; err == nil && streamErr != nil {
			err = streamErr
		}
		return result, err
	}

	return provider.ForwardStreamHTTP(ctx, s.streamClient.Client, "POST", url, apiKey, p.Type, actualBody, w, flusher)
}

type forwardParams struct {
	ingress  *model.IngressKey
	endpoint provider.EndpointType
	category model.ProviderCategory
	body     []byte
	stream   bool
}

func (s *Service) prepareForward(params forwardParams) (modelName string, providers []model.Provider, route *routeTarget, err error) {
	if len(params.body) > 0 {
		var payload map[string]any
		if err := json.Unmarshal(params.body, &payload); err != nil {
			return "", nil, nil, fmt.Errorf("%w: %v", ErrInvalidJSON, err)
		}
		modelName, _ = payload["model"].(string)
	}
	if params.endpoint != provider.EndpointListModels && modelName == "" {
		return "", nil, nil, ErrModelRequired
	}
	allowedProviders := cache.ParseJSONList(params.ingress.AllowedProvidersJSON)
	if modelName != "" {
		allowedModels := cache.ParseJSONList(params.ingress.AllowedModelsJSON)
		if !cache.MatchesList(allowedModels, modelName) && !cache.MatchesList(allowedModels, "*") {
			return "", nil, nil, ErrModelNotAllowed
		}
	}
	if !limiter.CheckIngressBudget(params.ingress, 0) {
		return "", nil, nil, ErrBudgetExceeded
	}
	if !limiter.CheckRateLimit(s.cache, params.ingress.ID, params.ingress.RateLimitRPM, params.ingress.RateLimitTPM, 1000) {
		return "", nil, nil, ErrRateLimitExceeded
	}
	route = s.resolveRoute(params.ingress.ID, modelName)
	providers, err = s.findProvidersForModel(modelName, allowedProviders, params.category, route)
	return modelName, providers, route, err
}

func (s *Service) recordUsage(ingress *model.IngressKey, att *attempt, clientIP, modelName, endpoint string, stream bool, tokensIn, tokensOut int, latency int64, status, errMsg, requestBody, responseBody string) {
	var providerKeyID, providerID uint
	if att != nil {
		providerKeyID = att.key.ID
		providerID = att.provider.ID
	}
	cost := pricing.Estimate(modelName, tokensIn, tokensOut, endpoint)
	s.cache.RecordUsage(cache.UsageDelta{
		IngressKeyID:  ingress.ID,
		ProviderKeyID: providerKeyID,
		ProviderID:    providerID,
		ClientIP:      clientIP,
		Endpoint:      endpoint,
		Model:         modelName,
		TokensIn:      tokensIn,
		TokensOut:     tokensOut,
		Cost:          cost,
		LatencyMs:     latency,
		Status:        status,
		Stream:        stream,
		ErrorMsg:      errMsg,
		RequestBody:   requestBody,
		ResponseBody:  responseBody,
		BudgetUsed:    cost,
		Requests:      1,
		Tokens:        int64(tokensIn + tokensOut),
	})
}

func (s *Service) ForwardEndpoint(ctx context.Context, ingress *model.IngressKey, endpoint provider.EndpointType, rawBody []byte, clientIP string) (*provider.ChatResponse, error) {
	category := categoryForEndpoint(endpoint)
	modelName, providers, route, err := s.prepareForward(forwardParams{
		ingress: ingress, endpoint: endpoint, category: category, body: rawBody,
	})
	if err != nil {
		s.recordUsage(ingress, nil, clientIP, modelNameOrBody(rawBody, modelName), string(endpoint), false, 0, 0, 0, "error", err.Error(), truncateLogBody(rawBody), "")
		return nil, err
	}
	attempts, err := s.buildAttempts(providers, modelName, route)
	if err != nil {
		s.recordUsage(ingress, nil, clientIP, modelName, string(endpoint), false, 0, 0, 0, "error", err.Error(), truncateLogBody(rawBody), "")
		return nil, err
	}
	start := time.Now()
	resp, att, err := s.forwardWithFailover(ctx, attempts, endpoint, rawBody)
	latency := time.Since(start).Milliseconds()
	status := "success"
	if err != nil || (resp != nil && resp.StatusCode >= 400) {
		status = "error"
	}
	errMsg := describeForwardError(err, resp)
	tokensIn, tokensOut := 0, 0
	var respBody []byte
	if resp != nil {
		tokensIn, tokensOut = resp.TokensIn, resp.TokensOut
		respBody = resp.Body
	}
	s.recordUsage(ingress, att, clientIP, modelName, string(endpoint), false, tokensIn, tokensOut, latency, status, errMsg, truncateLogBody(rawBody), truncateLogBody(respBody))
	if err != nil && resp != nil && resp.StatusCode >= 400 {
		return resp, nil
	}
	return resp, err
}

func modelNameOrBody(body []byte, modelName string) string {
	if modelName != "" {
		return modelName
	}
	return extractModelName(body)
}

func (s *Service) ForwardStream(ctx context.Context, ingress *model.IngressKey, endpoint provider.EndpointType, rawBody []byte, clientIP string, w io.Writer, flusher http.Flusher) error {
	category := categoryForEndpoint(endpoint)
	modelName, providers, route, err := s.prepareForward(forwardParams{
		ingress: ingress, endpoint: endpoint, category: category, body: rawBody, stream: true,
	})
	if err != nil {
		s.recordUsage(ingress, nil, clientIP, modelNameOrBody(rawBody, modelName), string(endpoint), true, 0, 0, 0, "error", err.Error(), truncateLogBody(rawBody), "")
		return err
	}
	attempts, err := s.buildAttempts(providers, modelName, route)
	if err != nil {
		s.recordUsage(ingress, nil, clientIP, modelName, string(endpoint), true, 0, 0, 0, "error", err.Error(), truncateLogBody(rawBody), "")
		return err
	}
	var lastErr error
	var lastAtt *attempt
	maxRetries := s.settings.GatewayMaxRetries()
	if maxRetries < 1 {
		maxRetries = 1
	}
	backoffMs := s.settings.GatewayRetryBackoffMs()
	for _, att := range attempts {
		if !s.canServeEndpoint(att.provider, endpoint) {
			continue
		}
		for try := 0; try < maxRetries; try++ {
			if try > 0 {
				time.Sleep(time.Duration(backoffMs) * time.Millisecond)
			}
			start := time.Now()
			result, err := s.forwardStream(ctx, att.provider, att.apiKey, endpoint, rawBody, w, flusher)
			latency := time.Since(start).Milliseconds()
			lastAtt = &att
			if err != nil {
				lastErr = err
				continue
			}
			if result == nil {
				lastErr = fmt.Errorf("empty upstream stream response")
				continue
			}
			if result.StatusCode < 400 {
				s.recordUsage(ingress, &att, clientIP, modelName, string(endpoint), true, result.TokensIn, result.TokensOut, latency, "success", "", truncateLogBody(rawBody), "[SSE stream]")
				return nil
			}
			errMsg := describeForwardError(nil, &provider.ChatResponse{StatusCode: result.StatusCode, Body: result.Body})
			lastErr = fmt.Errorf("%s", errMsg)
			respBody := "[SSE stream]"
			if len(result.Body) > 0 {
				respBody = truncateLogBody(result.Body)
			}
			if passThroughUpstreamStatus(result.StatusCode) {
				s.recordUsage(ingress, &att, clientIP, modelName, string(endpoint), true, result.TokensIn, result.TokensOut, latency, "error", errMsg, truncateLogBody(rawBody), respBody)
				return nil
			}
			if tryNextKeyOnUpstreamStatus(result.StatusCode) {
				break
			}
			if retryableUpstreamStatus(result.StatusCode) {
				continue
			}
			s.recordUsage(ingress, &att, clientIP, modelName, string(endpoint), true, result.TokensIn, result.TokensOut, latency, "error", errMsg, truncateLogBody(rawBody), respBody)
			return nil
		}
	}
	if lastErr == nil {
		lastErr = fmt.Errorf("all providers failed")
	}
	s.recordUsage(ingress, lastAtt, clientIP, modelName, string(endpoint), true, 0, 0, 0, "error", lastErr.Error(), truncateLogBody(rawBody), "[SSE stream]")
	return lastErr
}

func categoryForEndpoint(ep provider.EndpointType) model.ProviderCategory {
	switch ep {
	case provider.EndpointEmbeddings:
		return model.CategoryEmbedding
	case provider.EndpointRerank:
		return model.CategoryRerank
	default:
		return model.CategoryLLM
	}
}

func listModelsCategories() []model.ProviderCategory {
	return []model.ProviderCategory{
		model.CategoryLLM,
		model.CategoryEmbedding,
		model.CategoryRerank,
	}
}

func providerInListModelsCategory(cat model.ProviderCategory) bool {
	for _, c := range listModelsCategories() {
		if c == cat {
			return true
		}
	}
	return false
}

func (s *Service) ListModels(ctx context.Context, ingress *model.IngressKey) ([]map[string]any, error) {
	allowedProviders := cache.ParseJSONList(ingress.AllowedProvidersJSON)
	var result []map[string]any
	var hasLLM bool
	for _, p := range s.cache.ListProviders() {
		if !providerInListModelsCategory(p.Category) {
			continue
		}
		if !cache.MatchesList(allowedProviders, p.Name) && !cache.MatchesList(allowedProviders, string(p.Type)) {
			continue
		}
		if !s.providerAllowsEndpoint(&p, provider.EndpointListModels) {
			continue
		}
		if p.Category == model.CategoryLLM {
			hasLLM = true
		}
		for _, m := range s.cache.GetModels(p.ID) {
			if m.Enabled {
				result = append(result, map[string]any{
					"id": IngressModelID(m), "object": "model", "owned_by": p.Name,
				})
			}
		}
	}
	if len(result) == 0 && hasLLM {
		result = append(result, map[string]any{"id": "gpt-4o-mini", "object": "model", "owned_by": "luminary"})
	}
	return result, nil
}

func (s *Service) FetchModels(ctx context.Context, providerID uint) ([]provider.ModelInfo, error) {
	p, ok := s.cache.GetProvider(providerID)
	if !ok {
		return nil, fmt.Errorf("provider not found")
	}
	if !s.providerAllowsEndpoint(&p, provider.EndpointListModels) {
		return nil, fmt.Errorf("list_models endpoint not enabled")
	}
	keys := s.cache.ListProviderKeys(providerID)
	var firstKey string
	for _, k := range keys {
		if k.Enabled {
			plain, err := s.DecryptAPIKey(k.APIKeyEnc)
			if err == nil {
				firstKey = plain
				break
			}
		}
	}
	if firstKey == "" && p.Type != model.ProviderOllama {
		return nil, fmt.Errorf("no enabled keys")
	}
	cfg := s.providerConfig(&p)
	if p.Type == model.ProviderAnthropic {
		return s.anthropic.ListModelsWithConfig(ctx, firstKey, cfg)
	}
	return s.openai.ListModelsWithConfig(ctx, firstKey, cfg)
}

func (s *Service) HealthCheckProvider(ctx context.Context, providerID uint) (model.ProviderStatus, int64, string) {
	p, ok := s.cache.GetProvider(providerID)
	if !ok {
		return model.ProviderUnhealthy, 0, "not found"
	}
	if !model.IsCategorySupported(p.Category) {
		return model.ProviderUnhealthy, 0, "category not supported"
	}
	if p.Type == model.ProviderComfyUI {
		start := time.Now()
		if err := comfyui.New(p.BaseURL).HealthCheck(ctx); err != nil {
			return model.ProviderUnhealthy, time.Since(start).Milliseconds(), err.Error()
		}
		return model.ProviderHealthy, time.Since(start).Milliseconds(), ""
	}
	keys := s.cache.ListProviderKeys(providerID)
	if len(keys) == 0 && p.Type != model.ProviderOllama {
		return model.ProviderUnhealthy, 0, "no provider keys configured"
	}
	cfg := s.providerConfig(&p)
	return s.runHealthChecks(ctx, &p, keys, cfg)
}