package main

import (
	"crypto/sha256"
	"crypto/subtle"
	"encoding/hex"
	"encoding/json"
	"os"
	"sync"
	"time"
)

// API keys for the gateway. Only a SHA-256 of each key is stored; the key itself is shown once at creation.

const (
	keyMaxConcurrent = 4  // simultaneous requests per key
	keyRPM           = 60 // requests per minute per key
)

type APIKey struct {
	ID       string `json:"id"`
	Name     string `json:"name"`
	Hash     string `json:"hash"`
	Prefix   string `json:"prefix"` // first characters, for recognising a key in the list
	Created  int64  `json:"created"`
	Disabled bool   `json:"disabled"`
	Scope    string `json:"scope,omitempty"` // "" = gateway/jobs; "talk" = messages; "talk-worker" = own task inbox + messages
}

type KeyUsage struct {
	Requests         int64 `json:"requests"`
	Errors           int64 `json:"errors"`
	PromptTokens     int64 `json:"prompt_tokens"`
	CompletionTokens int64 `json:"completion_tokens"`
	LastUsed         int64 `json:"last_used"`
}

type keyFile struct {
	Keys  []APIKey             `json:"keys"`
	Usage map[string]*KeyUsage `json:"usage"`
}

type KeyStore struct {
	mu    sync.Mutex
	path  string
	keys  []APIKey
	usage map[string]*KeyUsage
	dirty bool
	conc  map[string]int
	win   map[string][]time.Time
}

func newKeyStore(path string) *KeyStore {
	ks := &KeyStore{path: path, usage: map[string]*KeyUsage{}, conc: map[string]int{}, win: map[string][]time.Time{}}
	if b, err := os.ReadFile(path); err == nil {
		var f keyFile
		if json.Unmarshal(b, &f) == nil {
			ks.keys = f.Keys
			if f.Usage != nil {
				ks.usage = f.Usage
			}
		}
	}
	return ks
}

func (ks *KeyStore) saveLocked() error {
	b, err := json.MarshalIndent(keyFile{Keys: ks.keys, Usage: ks.usage}, "", "  ")
	if err != nil {
		return err
	}
	tmp := ks.path + ".tmp"
	if err := os.WriteFile(tmp, b, 0600); err != nil {
		return err
	}
	if err := os.Rename(tmp, ks.path); err != nil {
		return err
	}
	ks.dirty = false
	return nil
}

// Pairing must not return a credential that would disappear on a server restart.
func (ks *KeyStore) createWorker(name string) (string, APIKey, error) {
	key := "wxk_" + randHex(24)
	k := APIKey{ID: randHex(8), Name: name, Hash: hashKey(key), Prefix: key[:8], Created: time.Now().Unix(), Scope: "talk-worker"}
	ks.mu.Lock()
	defer ks.mu.Unlock()
	ks.keys = append(ks.keys, k)
	ks.usage[k.ID] = &KeyUsage{}
	if err := ks.saveLocked(); err != nil {
		ks.keys = ks.keys[:len(ks.keys)-1]
		delete(ks.usage, k.ID)
		return "", APIKey{}, err
	}
	return key, k, nil
}

// flushLoop persists usage counters now and then (cheap, and at most 30 s of counts are at risk).
func (ks *KeyStore) flushLoop() {
	for range time.Tick(30 * time.Second) {
		ks.mu.Lock()
		if ks.dirty {
			ks.saveLocked()
		}
		ks.mu.Unlock()
	}
}

func hashKey(k string) string {
	h := sha256.Sum256([]byte(k))
	return hex.EncodeToString(h[:])
}

// create returns the new key in clear text; this is the only time it is available.
func (ks *KeyStore) create(name string) (string, APIKey) { return ks.createScoped(name, "") }

func (ks *KeyStore) createScoped(name, scope string) (string, APIKey) {
	key := "wxk_" + randHex(24)
	k := APIKey{ID: randHex(4), Name: name, Hash: hashKey(key), Prefix: key[:8], Created: time.Now().Unix(), Scope: scope}
	ks.mu.Lock()
	defer ks.mu.Unlock()
	ks.keys = append(ks.keys, k)
	ks.usage[k.ID] = &KeyUsage{}
	ks.saveLocked()
	return key, k
}

// authenticate is for the gateway and the job API: a key made for another scope does not open them.
func (ks *KeyStore) authenticate(token string) (APIKey, bool) { return ks.authenticateScope(token, "") }

func (ks *KeyStore) authenticateScope(token, scope string) (APIKey, bool) {
	if token == "" {
		return APIKey{}, false
	}
	h := hashKey(token)
	ks.mu.Lock()
	defer ks.mu.Unlock()
	for _, k := range ks.keys {
		if subtle.ConstantTimeCompare([]byte(k.Hash), []byte(h)) == 1 {
			return k, !k.Disabled && k.Scope == scope
		}
	}
	return APIKey{}, false
}

func (ks *KeyStore) setDisabled(id string, off bool) bool {
	ks.mu.Lock()
	defer ks.mu.Unlock()
	for i := range ks.keys {
		if ks.keys[i].ID == id {
			ks.keys[i].Disabled = off
			ks.saveLocked()
			return true
		}
	}
	return false
}

func (ks *KeyStore) remove(id string) bool {
	ks.mu.Lock()
	defer ks.mu.Unlock()
	for i := range ks.keys {
		if ks.keys[i].ID == id {
			ks.keys = append(ks.keys[:i], ks.keys[i+1:]...)
			delete(ks.usage, id)
			ks.saveLocked()
			return true
		}
	}
	return false
}

type keyView struct {
	APIKey
	Usage KeyUsage `json:"usage"`
}

func (ks *KeyStore) list() []keyView {
	ks.mu.Lock()
	defer ks.mu.Unlock()
	out := []keyView{}
	for _, k := range ks.keys {
		v := keyView{APIKey: k}
		if u := ks.usage[k.ID]; u != nil {
			v.Usage = *u
		}
		v.Hash = ""
		out = append(out, v)
	}
	return out
}

// admit enforces the per-key limits. The returned func must be called when the request ends.
func (ks *KeyStore) admit(id string) (func(), string) {
	ks.mu.Lock()
	defer ks.mu.Unlock()
	now := time.Now()
	w := ks.win[id][:0]
	for _, t := range ks.win[id] {
		if now.Sub(t) < time.Minute {
			w = append(w, t)
		}
	}
	if len(w) >= keyRPM {
		ks.win[id] = w
		return nil, "rate limit: too many requests per minute"
	}
	if ks.conc[id] >= keyMaxConcurrent {
		ks.win[id] = w
		return nil, "rate limit: too many simultaneous requests"
	}
	ks.win[id] = append(w, now)
	ks.conc[id]++
	return func() {
		ks.mu.Lock()
		ks.conc[id]--
		ks.mu.Unlock()
	}, ""
}

func (ks *KeyStore) record(id string, promptTok, complTok int64, failed bool) {
	ks.mu.Lock()
	defer ks.mu.Unlock()
	u := ks.usage[id]
	if u == nil {
		u = &KeyUsage{}
		ks.usage[id] = u
	}
	u.Requests++
	if failed {
		u.Errors++
	}
	u.PromptTokens += promptTok
	u.CompletionTokens += complTok
	u.LastUsed = time.Now().Unix()
	ks.dirty = true
}