package artifact

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"
)

// imageUsage mirrors the usage object returned by GPT Image endpoints. Keeping
// it typed avoids losing precision when provider metadata is persisted.
type imageUsage struct {
	TotalTokens      int64 `json:"total_tokens"`
	InputTokens      int64 `json:"input_tokens"`
	OutputTokens     int64 `json:"output_tokens"`
	TextInputTokens  int64 `json:"text_input_tokens"`
	ImageInputTokens int64 `json:"image_input_tokens"`
}

func (u imageUsage) available() bool {
	return u.TotalTokens > 0 || u.InputTokens > 0 || u.OutputTokens > 0 || u.TextInputTokens > 0 || u.ImageInputTokens > 0
}

type imagePricing struct {
	TextInputPerMTok   float64
	ImageInputPerMTok  float64
	ImageOutputPerMTok float64
	Basis              string
}

// openAIImagePricing intentionally supports only models whose public standard
// token rates are pinned here. Unknown/future model names still get full usage
// logging, but their local cost remains NULL instead of inventing a price.
func openAIImagePricing(model string) (imagePricing, bool) {
	m := strings.ToLower(strings.TrimSpace(model))
	switch {
	case m == "gpt-image-2" || strings.HasPrefix(m, "gpt-image-2-"):
		return imagePricing{
			TextInputPerMTok:   5.00,
			ImageInputPerMTok:  8.00,
			ImageOutputPerMTok: 30.00,
			Basis:              "openai-standard-2026-08-10:gpt-image-2:text-in-5,image-in-8,image-out-30-usd-per-mtok",
		}, true
	case m == "gpt-image-1.5" || strings.HasPrefix(m, "gpt-image-1.5-"):
		return imagePricing{
			TextInputPerMTok:   5.00,
			ImageInputPerMTok:  8.00,
			ImageOutputPerMTok: 32.00,
			Basis:              "openai-standard-2026-08-10:gpt-image-1.5:text-in-5,image-in-8,image-out-32-usd-per-mtok",
		}, true
	case m == "chatgpt-image-latest":
		return imagePricing{
			TextInputPerMTok:   5.00,
			ImageInputPerMTok:  8.00,
			ImageOutputPerMTok: 32.00,
			Basis:              "openai-standard-2026-08-10:chatgpt-image-latest:text-in-5,image-in-8,image-out-32-usd-per-mtok",
		}, true
	case m == "gpt-image-1":
		return imagePricing{
			TextInputPerMTok:   5.00,
			ImageInputPerMTok:  10.00,
			ImageOutputPerMTok: 40.00,
			Basis:              "openai-standard-2026-08-10:gpt-image-1:text-in-5,image-in-10,image-out-40-usd-per-mtok",
		}, true
	case m == "gpt-image-1-mini":
		return imagePricing{
			TextInputPerMTok:   2.00,
			ImageInputPerMTok:  2.50,
			ImageOutputPerMTok: 8.00,
			Basis:              "openai-standard-2026-08-10:gpt-image-1-mini:text-in-2,image-in-2.5,image-out-8-usd-per-mtok",
		}, true
	default:
		return imagePricing{}, false
	}
}

func estimateOpenAIImageCost(model string, u imageUsage) (float64, string, bool) {
	p, ok := openAIImagePricing(model)
	if !ok || !u.available() {
		return 0, "", false
	}
	// GPT Image reports the text/image input split. It matters because the two
	// modalities have different rates. If the split is missing, store usage but
	// do not claim a precise local cost.
	if u.InputTokens > 0 && u.TextInputTokens+u.ImageInputTokens <= 0 {
		return 0, p.Basis, false
	}
	if split := u.TextInputTokens + u.ImageInputTokens; u.InputTokens > 0 && split > 0 && split != u.InputTokens {
		return 0, p.Basis, false
	}
	cost := (float64(u.TextInputTokens)*p.TextInputPerMTok +
		float64(u.ImageInputTokens)*p.ImageInputPerMTok +
		float64(u.OutputTokens)*p.ImageOutputPerMTok) / 1_000_000
	return cost, p.Basis, true
}

func (w *Worker) recordOpenAIUsage(ctx context.Context, taskID, kind string, res imageResult) error {
	// Log every successful OpenAI image call, even if a provider response omits
	// the optional usage block. That keeps the call/card counters truthful while
	// leaving the local cost NULL when there is not enough data to calculate it.
	u := imageUsage{}
	if res.Usage != nil {
		u = *res.Usage
	}
	model := metaString(res.Meta, "model")
	endpoint := metaString(res.Meta, "endpoint")
	size := metaString(res.Meta, "size")
	quality := metaString(res.Meta, "quality")
	requestID := metaString(res.Meta, "request_id")

	var cost any
	if res.EstimatedCostUSD != nil {
		cost = *res.EstimatedCostUSD
	}
	metaJSON, _ := json.Marshal(res.Meta)
	_, err := w.db.ExecContext(ctx, `INSERT INTO artifact_api_usage(
		created_at,task_id,kind,provider,model,endpoint,size,quality,request_id,
		input_tokens,input_text_tokens,input_image_tokens,output_tokens,total_tokens,
		estimated_cost_usd,pricing_basis,meta_json)
		VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)`,
		time.Now().UTC().UnixMilli(), nullableString(taskID), kind, res.Provider, model, endpoint, size, quality, nullableString(requestID),
		u.InputTokens, u.TextInputTokens, u.ImageInputTokens, u.OutputTokens, u.TotalTokens,
		cost, res.PricingBasis, string(metaJSON))
	return err
}

func nullableString(s string) any {
	if strings.TrimSpace(s) == "" {
		return nil
	}
	return s
}

func metaString(meta map[string]any, key string) string {
	if meta == nil {
		return ""
	}
	v, ok := meta[key]
	if !ok || v == nil {
		return ""
	}
	if s, ok := v.(string); ok {
		return s
	}
	return fmt.Sprint(v)
}