package aiaudit

import (
	"context"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgxpool"
)

// Cost accounting. Kept apart from the prompt capture above on purpose:
// ai_call_logs is a 7-day debugging buffer, while spend has to answer "what did
// this tenant cost me last quarter" for years. Usage is therefore rolled up per
// day / company / user / model / role as each call completes.
//
// Money is micro-USD integers end to end. Summing millions of fractional-cent
// calls in float drifts; integers add up exactly.

// UsageRetentionDays bounds the rollup. Three years plus a leap day.
const UsageRetentionDays = 3*365 + 1

// MicrosPerUSD converts the stored integer unit to dollars.
const MicrosPerUSD = 1_000_000

// ModelPrice is micro-USD per 1,000,000 tokens.
type ModelPrice struct {
	InputPerMTok       int64
	CachedInputPerMTok int64
	OutputPerMTok      int64
}

// defaultModelPrices is the fallback when ai_model_prices has no row for a model
// (a provider the admin pointed at without adding a price). Keeping the shipped
// GPT-5.6 rates here means a fresh database still costs correctly before the
// migration seed is edited. Rates: OpenAI list price, 2026-07-30.
var defaultModelPrices = map[string]ModelPrice{
	"gpt-5.6-luna":  {InputPerMTok: 200_000, CachedInputPerMTok: 20_000, OutputPerMTok: 1_200_000},
	"gpt-5.6-terra": {InputPerMTok: 2_000_000, CachedInputPerMTok: 200_000, OutputPerMTok: 12_000_000},
	"gpt-5.6-sol":   {InputPerMTok: 5_000_000, CachedInputPerMTok: 500_000, OutputPerMTok: 30_000_000},
}

// DefaultModelPrice returns the shipped price for a model, ok=false when unknown.
func DefaultModelPrice(model string) (ModelPrice, bool) {
	p, ok := defaultModelPrices[NormalizeModel(model)]
	return p, ok
}

// NormalizeModel folds provider prefixes and dated suffixes onto the price key, so
// "openai/gpt-5.6-luna" and "gpt-5.6-luna-2026-07-30" both price as gpt-5.6-luna.
func NormalizeModel(model string) string {
	m := strings.ToLower(strings.TrimSpace(model))
	if i := strings.LastIndex(m, "/"); i >= 0 {
		m = m[i+1:]
	}
	// Trim a trailing -YYYY-MM-DD snapshot tag.
	if len(m) > 11 {
		tail := m[len(m)-11:]
		if tail[0] == '-' && isDigits(tail[1:5]) && tail[5] == '-' && isDigits(tail[6:8]) && tail[8] == '-' && isDigits(tail[9:11]) {
			m = m[:len(m)-11]
		}
	}
	return m
}

func isDigits(s string) bool {
	for _, r := range s {
		if r < '0' || r > '9' {
			return false
		}
	}
	return len(s) > 0
}

// CostMicros prices one call. Cached prompt tokens are billed at the cached rate
// and subtracted from fresh input — they are a subset of prompt tokens, not extra.
func CostMicros(p ModelPrice, promptTokens, cachedPromptTokens, completionTokens int) int64 {
	if cachedPromptTokens < 0 {
		cachedPromptTokens = 0
	}
	if cachedPromptTokens > promptTokens {
		cachedPromptTokens = promptTokens
	}
	fresh := int64(promptTokens - cachedPromptTokens)
	cached := int64(cachedPromptTokens)
	out := int64(completionTokens)
	if out < 0 {
		out = 0
	}
	if fresh < 0 {
		fresh = 0
	}
	// Rounded division keeps small calls from always costing zero.
	return divRound(fresh*p.InputPerMTok, 1_000_000) +
		divRound(cached*p.CachedInputPerMTok, 1_000_000) +
		divRound(out*p.OutputPerMTok, 1_000_000)
}

func divRound(n, d int64) int64 {
	if d == 0 {
		return 0
	}
	return (n + d/2) / d
}

// priceCache avoids a DB read per call; prices change at most a few times a year.
type priceCache struct {
	mu       sync.RWMutex
	prices   map[string]ModelPrice
	loadedAt time.Time
}

const priceCacheTTL = 5 * time.Minute

func (r *Recorder) priceFor(ctx context.Context, model string) ModelPrice {
	key := NormalizeModel(model)
	if r != nil {
		r.prices.mu.RLock()
		fresh := time.Since(r.prices.loadedAt) < priceCacheTTL
		p, ok := r.prices.prices[key]
		r.prices.mu.RUnlock()
		if fresh && ok {
			return p
		}
		if !fresh {
			r.reloadPrices(ctx)
			r.prices.mu.RLock()
			p, ok = r.prices.prices[key]
			r.prices.mu.RUnlock()
			if ok {
				return p
			}
		}
	}
	if p, ok := DefaultModelPrice(key); ok {
		return p
	}
	// Unknown model: record usage with zero cost rather than guessing a rate.
	return ModelPrice{}
}

func (r *Recorder) reloadPrices(ctx context.Context) {
	if r == nil || r.pool == nil {
		return
	}
	rows, err := r.pool.Query(ctx, `
		SELECT model, input_micros_per_mtok, cached_input_micros_per_mtok, output_micros_per_mtok
		FROM ai_model_prices`)
	if err != nil {
		// Keep whatever is cached; defaults still apply for misses.
		r.prices.mu.Lock()
		r.prices.loadedAt = time.Now()
		r.prices.mu.Unlock()
		return
	}
	defer rows.Close()
	loaded := map[string]ModelPrice{}
	for rows.Next() {
		var model string
		var in, cached, out int64
		if err := rows.Scan(&model, &in, &cached, &out); err != nil {
			continue
		}
		loaded[NormalizeModel(model)] = ModelPrice{
			InputPerMTok: in, CachedInputPerMTok: cached, OutputPerMTok: out,
		}
	}
	r.prices.mu.Lock()
	r.prices.prices = loaded
	r.prices.loadedAt = time.Now()
	r.prices.mu.Unlock()
}

// SourceLive marks rows written by the recorder from a provider usage block.
// SourceBackfill marks history reconstructed from processed_products (see
// BackfillFromProcessedProducts) — no user attribution, best effort.
const (
	SourceLive     = "live"
	SourceBackfill = "backfill"
)

const upsertUsageSQL = `
	INSERT INTO ai_usage_daily (
		day, company_id, user_id, model, role, source,
		calls, failed_calls, prompt_tokens, cached_prompt_tokens, completion_tokens,
		total_tokens, cost_micros, updated_at)
	VALUES (
		(now() AT TIME ZONE 'UTC')::date, $1, $2, $3, $4, 'live',
		1, $5, $6, $7, $8,
		$9, $10, now())
	ON CONFLICT (day, company_id, user_id, model, role, source) DO UPDATE SET
		calls                = ai_usage_daily.calls + 1,
		failed_calls         = ai_usage_daily.failed_calls + EXCLUDED.failed_calls,
		prompt_tokens        = ai_usage_daily.prompt_tokens + EXCLUDED.prompt_tokens,
		cached_prompt_tokens = ai_usage_daily.cached_prompt_tokens + EXCLUDED.cached_prompt_tokens,
		completion_tokens    = ai_usage_daily.completion_tokens + EXCLUDED.completion_tokens,
		total_tokens         = ai_usage_daily.total_tokens + EXCLUDED.total_tokens,
		cost_micros          = ai_usage_daily.cost_micros + EXCLUDED.cost_micros,
		updated_at           = now()`

// recordUsage rolls one call into the daily aggregate. Like Record it must never
// break the call it observes.
func (r *Recorder) recordUsage(ctx context.Context, c Call) {
	if r == nil || r.pool == nil {
		return
	}
	price := r.priceFor(ctx, c.Model)
	cost := CostMicros(price, c.PromptTokens, c.CachedPromptTokens, c.OutputTokens)
	failed := int64(0)
	if strings.TrimSpace(c.Error) != "" {
		failed = 1
	}
	writeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second)
	defer cancel()
	_, err := r.pool.Exec(writeCtx, upsertUsageSQL,
		orNilUUID(c.CompanyID), orNilUUID(c.UserID), NormalizeModel(c.Model), roleOrOther(c.Role),
		failed, c.PromptTokens, c.CachedPromptTokens, c.OutputTokens,
		c.TotalTokens, cost,
	)
	if err != nil {
		r.noteFailure(err)
	}
}

// orNilUUID keeps the primary key non-NULL (see the migration comment).
func orNilUUID(id uuid.UUID) uuid.UUID { return id }

func roleOrOther(role string) string {
	if r := strings.TrimSpace(role); r != "" {
		return r
	}
	return RoleOther
}

// CleanupExpiredUsage drops rollup rows past UsageRetentionDays.
func CleanupExpiredUsage(ctx context.Context, pool *pgxpool.Pool) (int64, error) {
	if pool == nil {
		return 0, nil
	}
	tag, err := pool.Exec(ctx, `
		DELETE FROM ai_usage_daily
		WHERE day < ((now() AT TIME ZONE 'UTC')::date - ($1::int))`, UsageRetentionDays)
	if err != nil {
		return 0, err
	}
	return tag.RowsAffected(), nil
}