// PG 实现：与内存实现同语义的直写 SQL。
package store

import (
	"database/sql"
	"time"
)

type pgStore struct {
	db *sql.DB
}

func (p *pgStore) InsertUsage(l UsageLog) error {
	_, err := p.db.Exec(
		`INSERT INTO usage_logs (id, run_id, character_id, user_id, model, slot, endpoint,
		 status, error_code, first_token_ms, total_ms, prompt_tokens, completion_tokens, at)
		 VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14)`,
		l.ID, l.RunID, l.CharacterID, l.UserID, l.Model, l.Slot, l.Endpoint,
		l.Status, l.ErrorCode, l.FirstTokenMS, l.TotalMS, l.PromptTokens, l.CompletionTokens, l.At,
	)
	return err
}

func (p *pgStore) AggregateUsage(userID string) ([]UsageAgg, error) {
	q := `SELECT user_id, model, endpoint, COUNT(*),
	      COUNT(*) FILTER (WHERE status = 'ok'),
	      COALESCE(SUM(prompt_tokens) FILTER (WHERE status = 'ok'), 0),
	      COALESCE(SUM(completion_tokens) FILTER (WHERE status = 'ok'), 0)
	      FROM usage_logs`
	args := []any{}
	if userID != "" {
		q += ` WHERE user_id = $1`
		args = append(args, userID)
	}
	q += ` GROUP BY user_id, model, endpoint ORDER BY user_id, model`
	rows, err := p.db.Query(q, args...)
	if err != nil {
		return nil, err
	}
	defer func() { _ = rows.Close() }()
	var out []UsageAgg
	for rows.Next() {
		var r UsageAgg
		if err := rows.Scan(&r.UserID, &r.Model, &r.Endpoint, &r.Requests, &r.OK, &r.PromptTokens, &r.CompletionTokens); err != nil {
			return nil, err
		}
		out = append(out, r)
	}
	return out, rows.Err()
}

func (p *pgStore) SampleUsage(limit int) ([]UsageLog, error) {
	rows, err := p.db.Query(
		`SELECT id, run_id, character_id, user_id, model, slot, endpoint, status, error_code,
		 first_token_ms, total_ms, prompt_tokens, completion_tokens, at
		 FROM usage_logs ORDER BY at DESC LIMIT $1`, limit)
	if err != nil {
		return nil, err
	}
	defer func() { _ = rows.Close() }()
	var out []UsageLog
	for rows.Next() {
		var l UsageLog
		if err := rows.Scan(&l.ID, &l.RunID, &l.CharacterID, &l.UserID, &l.Model, &l.Slot, &l.Endpoint,
			&l.Status, &l.ErrorCode, &l.FirstTokenMS, &l.TotalMS, &l.PromptTokens, &l.CompletionTokens, &l.At); err != nil {
			return nil, err
		}
		out = append(out, l)
	}
	return out, rows.Err()
}

func (p *pgStore) UpsertByok(k ByokKey) error {
	_, err := p.db.Exec(
		`INSERT INTO byok_keys (id, user_id, name, type, base_url, api_key, model, created_at)
		 VALUES ($1,$2,$3,$4,$5,$6,$7,$8)
		 ON CONFLICT (id) DO UPDATE SET name=$3, type=$4, base_url=$5, api_key=$6, model=$7`,
		k.ID, k.UserID, k.Name, k.Type, k.BaseURL, k.APIKey, k.Model, k.CreatedAt)
	return err
}

func (p *pgStore) ListByok(userID string) ([]ByokKey, error) {
	return p.byokQuery(`SELECT id, user_id, name, type, base_url, api_key, model, created_at, last_used, revoked_at
		FROM byok_keys WHERE user_id = $1 AND revoked_at IS NULL ORDER BY created_at`, userID)
}

func (p *pgStore) GetByok(userID, id string) (*ByokKey, error) {
	keys, err := p.byokQuery(`SELECT id, user_id, name, type, base_url, api_key, model, created_at, last_used, revoked_at
		FROM byok_keys WHERE user_id = $1 AND id = $2 AND revoked_at IS NULL`, userID, id)
	if err != nil || len(keys) == 0 {
		return nil, err
	}
	return &keys[0], nil
}

func (p *pgStore) byokQuery(q string, args ...any) ([]ByokKey, error) {
	rows, err := p.db.Query(q, args...)
	if err != nil {
		return nil, err
	}
	defer func() { _ = rows.Close() }()
	var out []ByokKey
	for rows.Next() {
		var k ByokKey
		if err := rows.Scan(&k.ID, &k.UserID, &k.Name, &k.Type, &k.BaseURL, &k.APIKey, &k.Model, &k.CreatedAt, &k.LastUsed, &k.RevokedAt); err != nil {
			return nil, err
		}
		out = append(out, k)
	}
	return out, rows.Err()
}

func (p *pgStore) RevokeByok(userID, id string) error {
	res, err := p.db.Exec(`UPDATE byok_keys SET revoked_at = $1 WHERE user_id = $2 AND id = $3 AND revoked_at IS NULL`,
		time.Now().UTC(), userID, id)
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return ErrByokNotFound
	}
	return nil
}

func (p *pgStore) TouchByok(id string, at time.Time) error {
	_, err := p.db.Exec(`UPDATE byok_keys SET last_used = $1 WHERE id = $2`, at, id)
	return err
}

func (p *pgStore) Reserve(userID string, amount int64) error {
	if amount <= 0 {
		return nil
	}
	// UPSERT + 原子累加；配额上限裁决在 quota 包（集中单一函数）
	_, err := p.db.Exec(`INSERT INTO quota_balances (user_id, reserved, used) VALUES ($1, $2, 0)
		ON CONFLICT (user_id) DO UPDATE SET reserved = quota_balances.reserved + $2`, userID, amount)
	return err
}

func (p *pgStore) Settle(userID string, reserved, actual int64) error {
	// 饱和截断在 SQL 层以 GREATEST 保证（防负）
	_, err := p.db.Exec(`UPDATE quota_balances SET
		reserved = GREATEST(reserved - $2, 0), used = used + $3 WHERE user_id = $1`,
		userID, reserved, actual)
	return err
}

func (p *pgStore) Refund(userID string, reserved int64) error {
	_, err := p.db.Exec(`UPDATE quota_balances SET reserved = GREATEST(reserved - $2, 0) WHERE user_id = $1`,
		userID, reserved)
	return err
}

func (p *pgStore) Balance(userID string) (QuotaBalance, error) {
	var b QuotaBalance
	b.UserID = userID
	err := p.db.QueryRow(`SELECT reserved, used FROM quota_balances WHERE user_id = $1`, userID).
		Scan(&b.Reserved, &b.Used)
	if err == sql.ErrNoRows {
		return b, nil
	}
	return b, err
}