GopherGate/internal/db/db.go

package db

import (
	"fmt"
	"log"
	"os"
	"path/filepath"
	"time"

	"github.com/jmoiron/sqlx"
	_ "modernc.org/sqlite"
	"golang.org/x/crypto/bcrypt"
)

type DB struct {
	*sqlx.DB
}

func Init(path string) (*DB, error) {
	// Ensure directory exists
	dir := filepath.Dir(path)
	if dir != "." && dir != "/" {
		if err := os.MkdirAll(dir, 0755); err != nil {
			return nil, fmt.Errorf("failed to create database directory: %w", err)
		}
	}

	// Connect to SQLite
	dsn := fmt.Sprintf("file:%s?_pragma=foreign_keys(1)", path)
	db, err := sqlx.Connect("sqlite", dsn)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to database: %w", err)
	}

	instance := &DB{db}

	// Run migrations
	if err := instance.RunMigrations(); err != nil {
		return nil, fmt.Errorf("failed to run migrations: %w", err)
	}

	return instance, nil
}

func (db *DB) RunMigrations() error {
	// Tables creation
	queries := []string{
		`CREATE TABLE IF NOT EXISTS clients (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            client_id TEXT UNIQUE NOT NULL,
            name TEXT,
            description TEXT,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            is_active BOOLEAN DEFAULT TRUE,
            rate_limit_per_minute INTEGER DEFAULT 60,
            total_requests INTEGER DEFAULT 0,
            total_tokens INTEGER DEFAULT 0,
            total_cost REAL DEFAULT 0.0
        )`,
		`CREATE TABLE IF NOT EXISTS llm_requests (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
            client_id TEXT,
            provider TEXT,
            model TEXT,
            prompt_tokens INTEGER,
            completion_tokens INTEGER,
            reasoning_tokens INTEGER DEFAULT 0,
            total_tokens INTEGER,
            cost REAL,
            has_images BOOLEAN DEFAULT FALSE,
            status TEXT DEFAULT 'success',
            error_message TEXT,
            duration_ms INTEGER,
            request_body TEXT,
            response_body TEXT,
            cache_read_tokens INTEGER DEFAULT 0,
            cache_write_tokens INTEGER DEFAULT 0,
            FOREIGN KEY (client_id) REFERENCES clients(client_id) ON DELETE SET NULL
        )`,
		`CREATE TABLE IF NOT EXISTS provider_configs (
            id TEXT PRIMARY KEY,
            display_name TEXT NOT NULL,
            enabled BOOLEAN DEFAULT TRUE,
            base_url TEXT,
            api_key TEXT,
            credit_balance REAL DEFAULT 0.0,
            low_credit_threshold REAL DEFAULT 5.0,
            billing_mode TEXT,
            api_key_encrypted BOOLEAN DEFAULT FALSE,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )`,
		`CREATE TABLE IF NOT EXISTS model_configs (
            id TEXT PRIMARY KEY,
            provider_id TEXT NOT NULL,
            display_name TEXT,
            enabled BOOLEAN DEFAULT TRUE,
            prompt_cost_per_m REAL,
            completion_cost_per_m REAL,
            cache_read_cost_per_m REAL,
            cache_write_cost_per_m REAL,
            mapping TEXT,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (provider_id) REFERENCES provider_configs(id) ON DELETE CASCADE
        )`,
		`CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT UNIQUE NOT NULL,
            password_hash TEXT NOT NULL,
            display_name TEXT,
            role TEXT DEFAULT 'admin',
            must_change_password BOOLEAN DEFAULT FALSE,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP
        )`,
		`CREATE TABLE IF NOT EXISTS client_tokens (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            client_id TEXT NOT NULL,
            token TEXT NOT NULL UNIQUE,
            name TEXT DEFAULT 'default',
            is_active BOOLEAN DEFAULT TRUE,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            last_used_at DATETIME,
            FOREIGN KEY (client_id) REFERENCES clients(client_id) ON DELETE CASCADE
        )`,
	}

	for _, q := range queries {
		if _, err := db.Exec(q); err != nil {
			return fmt.Errorf("migration failed for query [%s]: %w", q, err)
		}
	}

	// Add indices
	indices := []string{
		"CREATE INDEX IF NOT EXISTS idx_clients_client_id ON clients(client_id)",
		"CREATE INDEX IF NOT EXISTS idx_clients_created_at ON clients(created_at)",
		"CREATE INDEX IF NOT EXISTS idx_llm_requests_timestamp ON llm_requests(timestamp)",
		"CREATE INDEX IF NOT EXISTS idx_llm_requests_client_id ON llm_requests(client_id)",
		"CREATE INDEX IF NOT EXISTS idx_llm_requests_provider ON llm_requests(provider)",
		"CREATE INDEX IF NOT EXISTS idx_llm_requests_status ON llm_requests(status)",
		"CREATE UNIQUE INDEX IF NOT EXISTS idx_client_tokens_token ON client_tokens(token)",
		"CREATE INDEX IF NOT EXISTS idx_client_tokens_client_id ON client_tokens(client_id)",
		"CREATE INDEX IF NOT EXISTS idx_llm_requests_client_timestamp ON llm_requests(client_id, timestamp)",
		"CREATE INDEX IF NOT EXISTS idx_llm_requests_provider_timestamp ON llm_requests(provider, timestamp)",
		"CREATE INDEX IF NOT EXISTS idx_model_configs_provider_id ON model_configs(provider_id)",
	}

	for _, idx := range indices {
		if _, err := db.Exec(idx); err != nil {
			return fmt.Errorf("failed to create index [%s]: %w", idx, err)
		}
	}

	// Default admin user
	var count int
	if err := db.Get(&count, "SELECT COUNT(*) FROM users"); err != nil {
		return fmt.Errorf("failed to count users: %w", err)
	}

	if count == 0 {
		hash, err := bcrypt.GenerateFromPassword([]byte("admin123"), 12)
		if err != nil {
			return fmt.Errorf("failed to hash default password: %w", err)
		}
		_, err = db.Exec("INSERT INTO users (username, password_hash, role, must_change_password) VALUES ('admin', ?, 'admin', 1)", string(hash))
		if err != nil {
			return fmt.Errorf("failed to insert default admin: %w", err)
		}
		log.Println("Created default admin user with password 'admin123' (must change on first login)")
	}

	// Default client
	_, err := db.Exec(`INSERT OR IGNORE INTO clients (client_id, name, description)
        VALUES ('default', 'Default Client', 'Default client for anonymous requests')`)
	if err != nil {
		return fmt.Errorf("failed to insert default client: %w", err)
	}

	return nil
}

// Data models for DB tables

type Client struct {
	ID                  int       `db:"id"`
	ClientID            string    `db:"client_id"`
	Name                *string   `db:"name"`
	Description         *string   `db:"description"`
	CreatedAt           time.Time `db:"created_at"`
	UpdatedAt           time.Time `db:"updated_at"`
	IsActive            bool      `db:"is_active"`
	RateLimitPerMinute  int       `db:"rate_limit_per_minute"`
	TotalRequests       int       `db:"total_requests"`
	TotalTokens         int       `db:"total_tokens"`
	TotalCost           float64   `db:"total_cost"`
}

type LLMRequest struct {
	ID               int       `db:"id"`
	Timestamp        time.Time `db:"timestamp"`
	ClientID         *string   `db:"client_id"`
	Provider         *string   `db:"provider"`
	Model            *string   `db:"model"`
	PromptTokens     *int      `db:"prompt_tokens"`
	CompletionTokens *int      `db:"completion_tokens"`
	ReasoningTokens  int       `db:"reasoning_tokens"`
	TotalTokens      *int      `db:"total_tokens"`
	Cost             *float64  `db:"cost"`
	HasImages        bool      `db:"has_images"`
	Status           string    `db:"status"`
	ErrorMessage     *string   `db:"error_message"`
	DurationMS       *int      `db:"duration_ms"`
	RequestBody      *string   `db:"request_body"`
	ResponseBody     *string   `db:"response_body"`
	CacheReadTokens  int       `db:"cache_read_tokens"`
	CacheWriteTokens int       `db:"cache_write_tokens"`
}

type ProviderConfig struct {
	ID                 string    `db:"id"`
	DisplayName        string    `db:"display_name"`
	Enabled            bool      `db:"enabled"`
	BaseURL            *string   `db:"base_url"`
	APIKey             *string   `db:"api_key"`
	CreditBalance      float64   `db:"credit_balance"`
	LowCreditThreshold float64   `db:"low_credit_threshold"`
	BillingMode        *string   `db:"billing_mode"`
	APIKeyEncrypted    bool      `db:"api_key_encrypted"`
	UpdatedAt          time.Time `db:"updated_at"`
}

type ModelConfig struct {
	ID                  string    `db:"id"`
	ProviderID          string    `db:"provider_id"`
	DisplayName         *string   `db:"display_name"`
	Enabled             bool      `db:"enabled"`
	PromptCostPerM      *float64  `db:"prompt_cost_per_m"`
	CompletionCostPerM  *float64  `db:"completion_cost_per_m"`
	CacheReadCostPerM   *float64  `db:"cache_read_cost_per_m"`
	CacheWriteCostPerM  *float64  `db:"cache_write_cost_per_m"`
	Mapping             *string   `db:"mapping"`
	UpdatedAt           time.Time `db:"updated_at"`
}

type User struct {
	ID                 int       `db:"id"`
	Username           string    `db:"username"`
	PasswordHash       string    `db:"password_hash"`
	DisplayName        *string   `db:"display_name"`
	Role               string    `db:"role"`
	MustChangePassword bool      `db:"must_change_password"`
	CreatedAt          time.Time `db:"created_at"`
}

type ClientToken struct {
	ID        int        `db:"id"`
	ClientID  string     `db:"client_id"`
	Token     string     `db:"token"`
	Name      string     `db:"name"`
	IsActive  bool       `db:"is_active"`
	CreatedAt time.Time  `db:"created_at"`
	LastUsedAt *time.Time `db:"last_used_at"`
}