package gateway

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"
	"sync/atomic"
	"time"

	"llmsproxy/internal/config"
	"llmsproxy/internal/provider"
	"llmsproxy/internal/scheduler"
	"llmsproxy/internal/types"
)

// chatRequest mirrors the OpenAI chat completions request the gateway accepts.
type chatRequest struct {
	Model            string               `json:"model"`
	Messages         []types.ChatMessage  `json:"messages"`
	Temperature      *float64             `json:"temperature,omitempty"`
	MaxTokens        int                  `json:"max_tokens,omitempty"`
	Stream           bool                 `json:"stream,omitempty"`
	Tools            []interface{}        `json:"tools,omitempty"`
	ToolChoice       interface{}          `json:"tool_choice,omitempty"`
	DisableThinking  bool                 `json:"disable_thinking"`
	ExtraBody        map[string]interface{} `json:"extra_body,omitempty"`
}

// ChatCompletion is the non-streaming OpenAI response object.
type ChatCompletion struct {
	ID      string            `json:"id"`
	Object  string            `json:"object"`
	Created int64             `json:"created"`
	Model   string            `json:"model"`
	Choices []ChatChoice      `json:"choices"`
	Usage   *types.TokenUsage `json:"usage,omitempty"`
}

type ChatChoice struct {
	Index        int         `json:"index"`
	Message      RespMessage `json:"message"`
	FinishReason string      `json:"finish_reason"`
}

type RespMessage struct {
	Role             string          `json:"role,omitempty"`
	Content          string          `json:"content"`
	ReasoningContent string          `json:"reasoning_content,omitempty"`
	ToolCalls        json.RawMessage `json:"tool_calls,omitempty"`
}

type ChatChunk struct {
	ID      string        `json:"id"`
	Object  string        `json:"object"`
	Created int64         `json:"created"`
	Model   string        `json:"model"`
	Choices []ChunkChoice `json:"choices"`
}

type ChunkChoice struct {
	Index        int           `json:"index"`
	Delta        RespMessage   `json:"delta"`
	FinishReason *string       `json:"finish_reason"`
}

var seq int64

func newID() string {
	n := atomic.AddInt64(&seq, 1)
	return fmt.Sprintf("chatcmpl-%d", n)
}

func isAuto(m string) bool {
	m = strings.TrimSpace(m)
	return m == "" || strings.EqualFold(m, "AUTO")
}

// resolveCands picks the ordered candidate providers for a requested model.
// toolCalling requests are anchored: they resolve to exactly one provider
// (highest-priority available) so a tool-call round never switches models.
func (g *Gateway) resolveCands(ctx context.Context, req *chatRequest) ([]*provider.Provider, string) {
	model := req.Model
	if model == "" {
		model = g.core.DefaultModel()
	}
	cands, effective := g.resolveByModel(model)
	cands = chatOnly(cands)
	allow := g.allowedModels(ctx)
	if allow != nil {
		cands = filterCandsByModels(cands, allow)
	}
	if !toolRequest(req) {
		return cands, effective
	}
	// tool-call request: pin to one provider (no AUTO fallback across models)
	if len(cands) == 0 {
		return nil, effective
	}
	first := cands[0]
	eff := first.ModelFor(model)
	if eff == "" {
		eff = firstModel(first)
	}
	return []*provider.Provider{first}, eff
}

// filterCandsByModels keeps only providers exposing at least one model of the
// scope (used for user keys with a restricted model scope). An "AUTO" scope
// entry means the key is allowed to use any model. Scope entries with a
// Source pinned to a specific upstream narrow the candidates to that source
// for the matching model.
func filterCandsByModels(cands []*provider.Provider, allow []config.ModelScope) []*provider.Provider {
	for _, m := range allow {
		if m.Model == "" || strings.EqualFold(m.Model, "AUTO") {
			return cands
		}
	}
	allowed := make(map[string]bool, len(allow))
	byModelSrc := map[string]map[string]bool{}
	for _, m := range allow {
		allowed[m.Model] = true
		if m.Source != "" {
			if byModelSrc[m.Model] == nil {
				byModelSrc[m.Model] = map[string]bool{}
			}
			byModelSrc[m.Model][m.Source] = true
		}
	}
	out := make([]*provider.Provider, 0, len(cands))
	for _, p := range cands {
		for _, id := range p.Models() {
			if !allowed[id] {
				continue
			}
			if srcs := byModelSrc[id]; len(srcs) > 0 && !srcs[p.Name()] {
				continue
			}
			out = append(out, p)
			break
		}
	}
	return out
}

// intersectModels restricts a model list to the scope (preserving order). An
// "AUTO" scope entry grants every model.
func intersectModels(models []string, allow []config.ModelScope) []string {
	allowed := make(map[string]bool, len(allow))
	any := false
	for _, m := range allow {
		if m.Model == "" || strings.EqualFold(m.Model, "AUTO") {
			any = true
			break
		}
		allowed[m.Model] = true
	}
	if any {
		return models
	}
	out := make([]string, 0, len(models))
	seen := map[string]bool{}
	for _, m := range models {
		if allowed[m] && !seen[m] {
			seen[m] = true
			out = append(out, m)
		}
	}
	return out
}

// checkModelScope validates the effective model against the key's model scope
// and token quota. Returns an error message when rejected. A scope entry with
// model "AUTO" grants all models; its quota caps the key's total tokens.
func (g *Gateway) checkModelScope(ctx context.Context, model string) string {
	allow := g.allowedModels(ctx)
	if allow == nil {
		return ""
	}
	for _, sc := range allow {
		if sc.Model != "" && strings.EqualFold(sc.Model, "AUTO") {
			if sc.TokenQuota > 0 {
				used := g.scopeTokens(ctx, sc)
				if used >= sc.TokenQuota {
					return fmt.Sprintf("token quota exceeded (%d/%d)", used, sc.TokenQuota)
				}
			}
			return ""
		}
	}
	for _, sc := range allow {
		if sc.Model != model {
			continue
		}
		if sc.TokenQuota > 0 {
			used := g.scopeTokens(ctx, sc)
			if used >= sc.TokenQuota {
				return fmt.Sprintf("token quota exceeded for %q (%d/%d)", model, used, sc.TokenQuota)
			}
		}
		return ""
	}
	return fmt.Sprintf("model %q is not allowed for this key", model)
}

// scopeTokens returns the tokens a scope entry has consumed within its reset
// window (total for AUTO / per model otherwise).
func (g *Gateway) scopeTokens(ctx context.Context, sc config.ModelScope) int64 {
	k := keyID(reqKey(ctx))
	if sc.Model != "" && strings.EqualFold(sc.Model, "AUTO") {
		return g.stats.KeyTokens(k)
	}
	win := AutoPeriodSeconds(sc.Period, sc.Hours)
	return g.stats.WindowTokens(sc.Model, sc.Source, win)
}

func hasScopeModel(list []config.ModelScope, s string) bool {
	for _, x := range list {
		if x.Model == s || (x.Model != "" && strings.EqualFold(x.Model, "AUTO")) {
			return true
		}
	}
	return false
}

func (g *Gateway) resolveByModel(model string) ([]*provider.Provider, string) {
	if isAuto(model) {
		return g.core.Registry().Resolve("AUTO"), ""
	}
	return g.core.Registry().Resolve(model), model
}

// chatOnly keeps providers that expose at least one chat-capable model, so a
// chat/AUTO request never lands on an image-only source (or borrows its image
// model id). Explicit image-kind requests stay on the imageOnly path.
func chatOnly(cands []*provider.Provider) []*provider.Provider {
	out := make([]*provider.Provider, 0, len(cands))
	for _, p := range cands {
		for _, id := range p.Models() {
			if m := p.ModelByID(id); m == nil || m.Kind != "image" {
				out = append(out, p)
				break
			}
		}
	}
	return out
}

// toolRequest reports whether the request participates in a tool-call round.
func toolRequest(req *chatRequest) bool {
	if len(req.Tools) > 0 || req.ToolChoice != nil {
		return true
	}
	for _, m := range req.Messages {
		if m.Role == "tool" || len(m.ToolCalls) > 0 {
			return true
		}
	}
	return false
}

func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use POST")
		return
	}
	var req chatRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
		return
	}
	if len(req.Messages) == 0 {
		writeError(w, http.StatusBadRequest, "invalid_request", "messages is required")
		return
	}
	model := req.Model
	if model == "" {
		model = g.core.DefaultModel()
	}
	if isAuto(model) {
		plans := g.autoPlans()
		if len(plans) == 0 {
			writeError(w, http.StatusServiceUnavailable, "no_provider", "no auto slot available (quota exhausted or none configured)")
			return
		}
		if msg := g.checkModelScope(r.Context(), "AUTO"); msg != "" {
			writeError(w, http.StatusForbidden, "model_not_allowed", msg)
			return
		}
		ctx := r.Context()
		done := g.stats.Begin()
		defer done()
		inner := &types.ChatRequest{
			Messages:        req.Messages,
			Temperature:     req.Temperature,
			MaxTokens:       req.MaxTokens,
			Stream:          req.Stream,
			Tools:           req.Tools,
			ToolChoice:      req.ToolChoice,
			DisableThinking: req.DisableThinking,
			ExtraBody:       req.ExtraBody,
		}
		rec := &Req{
			Key:  keyID(reqKey(ctx)),
			Type: "chat",
			OK:   false,
		}
		if req.Stream {
			rec.Type = "stream"
			g.streamChatAuto(w, ctx, plans, inner, rec)
			return
		}
		g.singleChatAuto(w, ctx, plans, inner, rec)
		return
	}
	if !isAuto(model) {
		if allow := g.allowedModels(r.Context()); allow != nil && !hasScopeModel(allow, model) {
			writeError(w, http.StatusForbidden, "model_not_allowed", fmt.Sprintf("model %q is not allowed for this key", model))
			return
		}
	}
	cands, effective := g.resolveCands(r.Context(), &req)
	if len(cands) == 0 {
		writeError(w, http.StatusServiceUnavailable, "no_provider", "no LLM source configured")
		return
	}
	if effective == "" {
		effective = firstModel(cands[0])
	}
	if msg := g.checkModelScope(r.Context(), effective); msg != "" {
		writeError(w, http.StatusForbidden, "model_not_allowed", msg)
		return
	}
	ctx := r.Context()
	done := g.stats.Begin()
	defer done()

	inner := &types.ChatRequest{
		Model:           effective,
		Messages:        req.Messages,
		Temperature:     req.Temperature,
		MaxTokens:       req.MaxTokens,
		Stream:          req.Stream,
		Tools:           req.Tools,
		ToolChoice:      req.ToolChoice,
		DisableThinking: req.DisableThinking,
		ExtraBody:       req.ExtraBody,
	}
	rec := &Req{
		Key:    keyID(reqKey(ctx)),
		Type:   "chat",
		Model:  effective,
		Source: firstSource(cands),
		OK:     false,
	}
	if req.Stream {
		rec.Type = "stream"
		g.streamChat(w, ctx, cands, inner, effective, rec)
		return
	}
	g.singleChat(w, ctx, cands, inner, effective, rec)
}

func normalizeModel(m string) string {
	if isAuto(m) {
		return ""
	}
	return m
}

func firstModel(p *provider.Provider) string {
	ms := p.Models()
	if len(ms) > 0 {
		return ms[0]
	}
	return "auto"
}

// imageOnly keeps providers exposing at least one image-kind model.
func imageOnly(cands []*provider.Provider) []*provider.Provider {
	var out []*provider.Provider
	for _, p := range cands {
		for _, id := range p.Models() {
			if m := p.ModelByID(id); m != nil && m.Kind == "image" {
				out = append(out, p)
				break
			}
		}
	}
	return out
}

// toolCallsWire converts unified tool calls to the OpenAI wire format:
// tool_calls:[{id,type,function:{name,arguments:StringJSON}}]. Clients expect
// arguments to be a JSON string, not an object.
func toolCallsWire(tcs []types.ToolCall) json.RawMessage {
	wire := make([]map[string]interface{}, 0, len(tcs))
	for _, tc := range tcs {
		args := "{}"
		if tc.Arguments != nil {
			if b, err := json.Marshal(tc.Arguments); err == nil {
				args = string(b)
			}
		}
		wire = append(wire, map[string]interface{}{
			"id":   tc.ID,
			"type": tc.Type,
			"function": map[string]interface{}{
				"name":      tc.Name,
				"arguments": args,
			},
		})
	}
	b, _ := json.Marshal(wire)
	return b
}

func firstSource(cands []*provider.Provider) string {
	if len(cands) > 0 {
		return cands[0].Name()
	}
	return ""
}

// effectiveImageModel picks the image model id that will actually be used so
// quota checks can target it (AUTO resolves to the first image candidate).
func effectiveImageModel(model string, cands []*provider.Provider) string {
	if !isAuto(model) && model != "" {
		return model
	}
	if len(cands) > 0 {
		for _, id := range cands[0].Models() {
			if m := cands[0].ModelByID(id); m != nil && m.Kind == "image" {
				return id
			}
		}
	}
	return model
}

func estimatePromptTokens(req *types.ChatRequest) int64 {
	if req == nil {
		return 0
	}
	b, _ := json.Marshal(struct {
		Messages []types.ChatMessage `json:"messages"`
		Tools    []interface{}       `json:"tools,omitempty"`
	}{Messages: req.Messages, Tools: req.Tools})
	if len(b) == 0 {
		return 0
	}
	return int64(len(b)/3 + 1)
}

func estimateTextTokens(parts ...interface{}) int64 {
	var n int
	for _, p := range parts {
		switch v := p.(type) {
		case string:
			n += len(v)
		case json.RawMessage:
			n += len(v)
		case []types.ToolCall:
			b, _ := json.Marshal(v)
			n += len(b)
		}
	}
	if n == 0 {
		return 0
	}
	return int64(n/3 + 1)
}

func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
	rec.LatMs = 0
	t0 := time.Now()
	resp, usedSrc, usedModel, err := g.core.Scheduler().Chat(ctx, scheduler.FromRegistry(cands), req)
	rec.LatMs = time.Since(t0).Milliseconds()
	if err != nil {
		rec.OK = false
		rec.Status = http.StatusBadGateway
		rec.Err = err.Error()
		g.writeRec(rec)
		writeError(w, http.StatusBadGateway, "upstream_error", err.Error())
		return
	}
	rec.OK = true
	rec.Status = http.StatusOK
	rec.Prompt = int64(resp.TokenUsage.Prompt)
	if rec.Prompt == 0 {
		rec.Prompt = estimatePromptTokens(req)
	}
	rec.Compl = int64(resp.TokenUsage.Completion)
	if rec.Compl == 0 {
		rec.Compl = estimateTextTokens(resp.Content, resp.ReasoningContent, resp.ToolCalls)
	}
	rec.Source = usedSrc
	rec.Model = usedModel
	g.writeRec(rec)
	msg := RespMessage{Role: "assistant", Content: resp.Content}
	if resp.ReasoningContent != "" {
		msg.ReasoningContent = resp.ReasoningContent
	}
	if len(resp.ToolCalls) > 0 {
		msg.ToolCalls = toolCallsWire(resp.ToolCalls)
	}
	out := ChatCompletion{
		ID:      newID(),
		Object:  "chat.completion",
		Created: time.Now().Unix(),
		Model:   effective,
		Choices: []ChatChoice{{Index: 0, Message: msg, FinishReason: resp.FinishReason}},
	}
	if resp.TokenUsage.Total > 0 || resp.TokenUsage.Prompt > 0 || resp.TokenUsage.Completion > 0 {
		out.Usage = &resp.TokenUsage
	}
	writeJSON(w, http.StatusOK, out)
}

// writeRec records a finished request (audit + aggregates).
func (g *Gateway) writeRec(rec *Req) {
	if rec == nil {
		return
	}
	if rec.Time == 0 {
		rec.Time = time.Now().UnixMilli()
	}
	g.stats.Record(*rec)
}

func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
	rec.LatMs = 0
	t0 := time.Now()
	rec.OK = true
	rec.Status = http.StatusOK
	defer func() {
		rec.LatMs = time.Since(t0).Milliseconds()
		g.writeRec(rec)
	}()
	chunks, _, usedModel, err := g.core.Scheduler().ChatStream(ctx, scheduler.FromRegistry(cands), req)
	if err != nil {
		rec.OK = false
		rec.Status = http.StatusBadGateway
		rec.Err = err.Error()
		writeError(w, http.StatusBadGateway, "upstream_error", err.Error())
		return
	}
	if usedModel != "" {
		rec.Model = usedModel
	}
	rec.Prompt = estimatePromptTokens(req)
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.WriteHeader(http.StatusOK)

	flusher, _ := w.(http.Flusher)
	id := newID()
	created := time.Now().Unix()
	send := func(obj interface{}) bool {
		b, err := json.Marshal(obj)
		if err != nil {
			return false
		}
		if _, err := fmt.Fprintf(w, "data: %s\n\n", b); err != nil {
			return false
		}
		if flusher != nil {
			flusher.Flush()
		}
		return true
	}

	if !send(ChatChunk{
		ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
		Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{Role: "assistant"}}},
	}) {
		return
	}
	for ck := range chunks {
		chunk := ChatChunk{
			ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
		}
		delta := RespMessage{Role: "assistant", Content: ck.Content}
		if ck.ReasoningContent != "" {
			delta.ReasoningContent = ck.ReasoningContent
		}
		if len(ck.ToolCalls) > 0 {
			delta.ToolCalls = ck.ToolCalls
		}
		choice := ChunkChoice{Index: 0, Delta: delta}
		if ck.Done {
			stop := "stop"
			choice.FinishReason = &stop
		}
		chunk.Choices = []ChunkChoice{choice}
		rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)+len(ck.ToolCalls)) / 3
		if !send(chunk) {
			return
		}
	}
	stop := "stop"
	send(ChatChunk{
		ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
		Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
	})
	fmt.Fprintf(w, "data: [DONE]\n\n")
	if flusher != nil {
		flusher.Flush()
	}
}

// autoPlan is one schedulable AUTO slot: a model id pinned to its provider
// with an optional token quota window. Quota-exhausted slots are skipped.
type autoPlan struct {
	p     *provider.Provider
	model string
	quota int64
	win   int64
}

// autoPlans builds the schedulable AUTO slots from the persisted rules. A
// slot is schedulable while its model is available and (when quota > 0) the
// tokens used within its reset window are below the quota.
func (g *Gateway) autoPlans() []autoPlan {
	rules := g.core.AutoRules()
	if len(rules) == 0 {
		return nil
	}
	plans := make([]autoPlan, 0, len(rules))
	for _, e := range rules {
		p := g.core.ProviderForSlot(e.Model, e.Source)
		if p == nil {
			continue
		}
		if m := p.ModelByID(e.Model); m != nil && m.Kind == "image" {
			continue
		}
		win := AutoPeriodSeconds(e.Period, e.Hours)
		if e.TokenQuota > 0 && g.stats.WindowTokens(e.Model, e.Source, win) >= e.TokenQuota {
			continue
		}
		plans = append(plans, autoPlan{p: p, model: e.Model, quota: e.TokenQuota, win: win})
	}
	return plans
}

// singleChatAuto runs a non-streaming AUTO request slot by slot: each slot
// pins its own model; a slot whose provider errors out is skipped. The first
// slot to answer wins; when every slot fails, the recorded error is from the
// last one.
func (g *Gateway) singleChatAuto(w http.ResponseWriter, ctx context.Context, plans []autoPlan, req *types.ChatRequest, rec *Req) {
	rec.LatMs = 0
	t0 := time.Now()
	var lastErr error
	var lastSrc, lastModel string
	for _, pl := range plans {
		if !pl.p.Available() {
			continue
		}
		r := *req
		r.Model = pl.model
		lastSrc, lastModel = pl.p.Name(), pl.model
		resp, usedSrc, usedModel, err := g.core.Scheduler().Chat(ctx, scheduler.FromRegistry([]*provider.Provider{pl.p}), &r)
		if err != nil {
			lastErr = err
			continue
		}
		rec.LatMs = time.Since(t0).Milliseconds()
		rec.OK = true
		rec.Status = http.StatusOK
		rec.Prompt = int64(resp.TokenUsage.Prompt)
		if rec.Prompt == 0 {
			rec.Prompt = estimatePromptTokens(&r)
		}
		rec.Compl = int64(resp.TokenUsage.Completion)
		if rec.Compl == 0 {
			rec.Compl = estimateTextTokens(resp.Content, resp.ReasoningContent, resp.ToolCalls)
		}
		rec.Source = usedSrc
		rec.Model = usedModel
		g.writeRec(rec)
		msg := RespMessage{Role: "assistant", Content: resp.Content}
		if resp.ReasoningContent != "" {
			msg.ReasoningContent = resp.ReasoningContent
		}
		if len(resp.ToolCalls) > 0 {
			msg.ToolCalls = toolCallsWire(resp.ToolCalls)
		}
		out := ChatCompletion{
			ID:      newID(),
			Object:  "chat.completion",
			Created: time.Now().Unix(),
			Model:   usedModel,
			Choices: []ChatChoice{{Index: 0, Message: msg, FinishReason: resp.FinishReason}},
		}
		if resp.TokenUsage.Total > 0 || resp.TokenUsage.Prompt > 0 || resp.TokenUsage.Completion > 0 {
			out.Usage = &resp.TokenUsage
		}
		writeJSON(w, http.StatusOK, out)
		return
	}
	rec.LatMs = time.Since(t0).Milliseconds()
	if lastErr == nil {
		lastErr = fmt.Errorf("no provider available")
	}
	rec.OK = false
	rec.Status = http.StatusBadGateway
	rec.Err = lastErr.Error()
	if rec.Model == "" {
		rec.Model = lastModel
	}
	if rec.Source == "" {
		rec.Source = lastSrc
	}
	g.writeRec(rec)
	writeError(w, http.StatusBadGateway, "upstream_error", lastErr.Error())
}

// streamChatAuto streams an AUTO request. It stays pinned to the first slot
// whose stream begins; a slot that fails to connect is skipped.
func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, plans []autoPlan, req *types.ChatRequest, rec *Req) {
	rec.LatMs = 0
	t0 := time.Now()
	rec.OK = true
	rec.Status = http.StatusOK
	defer func() {
		rec.LatMs = time.Since(t0).Milliseconds()
		g.writeRec(rec)
	}()
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.WriteHeader(http.StatusOK)
	flusher, _ := w.(http.Flusher)
	id := newID()
	created := time.Now().Unix()
	send := func(obj interface{}) bool {
		b, err := json.Marshal(obj)
		if err != nil {
			return false
		}
		if _, err := fmt.Fprintf(w, "data: %s\n\n", b); err != nil {
			return false
		}
		if flusher != nil {
			flusher.Flush()
		}
		return true
	}
	if !send(ChatChunk{
		ID: id, Object: "chat.completion.chunk", Created: created, Model: "auto",
		Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{Role: "assistant"}}},
	}) {
		return
	}
	var lastErr error
	var lastSrc, lastModel string
	for _, pl := range plans {
		if !pl.p.Available() {
			continue
		}
		r := *req
		r.Model = pl.model
		lastSrc, lastModel = pl.p.Name(), pl.model
		chunks, usedSrc, usedModel, err := g.core.Scheduler().ChatStream(ctx, scheduler.FromRegistry([]*provider.Provider{pl.p}), &r)
		if err != nil {
			lastErr = err
			continue
		}
		if usedModel != "" {
			rec.Model = usedModel
			rec.Source = usedSrc
		}
		rec.Prompt = estimatePromptTokens(&r)
		for ck := range chunks {
			chunk := ChatChunk{
				ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
			}
			delta := RespMessage{Role: "assistant", Content: ck.Content}
			if ck.ReasoningContent != "" {
				delta.ReasoningContent = ck.ReasoningContent
			}
			if len(ck.ToolCalls) > 0 {
				delta.ToolCalls = ck.ToolCalls
			}
			choice := ChunkChoice{Index: 0, Delta: delta}
			if ck.Done {
				stop := "stop"
				choice.FinishReason = &stop
			}
			chunk.Choices = []ChunkChoice{choice}
			rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)+len(ck.ToolCalls)) / 3
			if !send(chunk) {
				return
			}
		}
		return
	}
	if lastErr == nil {
		lastErr = fmt.Errorf("no provider available")
	}
	rec.OK = false
	rec.Status = http.StatusBadGateway
	rec.Err = lastErr.Error()
	if rec.Model == "" {
		rec.Model = lastModel
	}
	if rec.Source == "" {
		rec.Source = lastSrc
	}
	errEvent, _ := json.Marshal(map[string]interface{}{"error": map[string]string{"message": lastErr.Error(), "type": "upstream_error"}})
	fmt.Fprintf(w, "data: %s\n\n", errEvent)
	fmt.Fprintf(w, "data: [DONE]\n\n")
	if flusher != nil {
		flusher.Flush()
	}
}

func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use POST")
		return
	}
	var req types.ImageGenRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
		return
	}
	if req.Prompt == "" {
		writeError(w, http.StatusBadRequest, "invalid_request", "prompt is required")
		return
	}
	model := req.Model
	if model == "" {
		model = g.core.DefaultModel()
	}
	if !isAuto(model) {
		if allow := g.allowedModels(r.Context()); allow != nil && !hasScopeModel(allow, model) {
			writeError(w, http.StatusForbidden, "model_not_allowed", fmt.Sprintf("model %q is not allowed for this key", model))
			return
		}
	}
	cands, _ := g.resolveByModel(model)
	cands = imageOnly(cands)
	if allow := g.allowedModels(r.Context()); allow != nil {
		cands = filterCandsByModels(cands, allow)
	}
	if len(cands) == 0 {
		writeError(w, http.StatusServiceUnavailable, "no_provider", "no image source configured")
		return
	}
	if msg := g.checkModelScope(r.Context(), effectiveImageModel(model, cands)); msg != "" {
		writeError(w, http.StatusForbidden, "model_not_allowed", msg)
		return
	}
	done := g.stats.Begin()
	defer done()
	rec := &Req{Key: keyID(reqKey(r.Context())), Type: "image", Model: model, Source: firstSource(cands), OK: false}
	t0 := time.Now()
	resp, usedSrc, err := g.core.Scheduler().Image(r.Context(), scheduler.FromRegistry(cands), &req)
	rec.LatMs = time.Since(t0).Milliseconds()
	if err != nil {
		rec.Status = http.StatusBadGateway
		rec.Err = err.Error()
		g.writeRec(rec)
		writeError(w, http.StatusBadGateway, "upstream_error", err.Error())
		return
	}
	if usedSrc != "" {
		rec.Source = usedSrc
	}
	rec.OK = true
	rec.Status = http.StatusOK
	rec.Compl = int64(len(resp.ImageData))
	g.writeRec(rec)
	writeJSON(w, http.StatusOK, types.ImageGenResponse{
		Created: time.Now().Unix(),
		Data:    resp.ImageData,
	})
}