zkill-susser/db.go

package main

import (
	"context"
	"fmt"
	"strconv"
	"strings"
	"time"

	"zkillsusser/models"

	logger "git.site.quack-lab.dev/dave/cylogger"
	"github.com/ClickHouse/clickhouse-go/v2"
	"github.com/ClickHouse/clickhouse-go/v2/lib/driver"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/schema"
)

type DB interface {
	Init() error
	Get() *gorm.DB

	SaveFlatKillmails(killmails []*FlatKillmail, attackers []FlatKillmailAttacker, items []FlatKillmailItem) error
	QueryFits(params QueryParams) (*FitStatistics, error)
	SearchShips(query string, limit int) ([]models.InvType, error)
	SearchSystems(query string, limit int) ([]models.MapSolarSystem, error)
	SearchModules(query string, limit int) ([]models.InvType, error)
	SearchGroups(query string, limit int) ([]models.InvGroup, error)

	// Non retarded APIs below
	GetItemTypes(itemIDs []int64) ([]models.InvType, error)
	GetSolarSystems(systemIDs []int64) ([]models.MapSolarSystem, error)
	ExpandGroupsIntoItemTypeIds(groups []int64) ([]int64, error)
	GetModuleSlots(moduleIDs []int64) (map[int64]ModuleSlot, error)
	CacheSet(key string, data []byte) error
	CacheGet(key string) ([]byte, error)
	CacheClean() error

	GetType(ctx context.Context, typeID int32) (*models.InvType, error)
	GetGroup(ctx context.Context, groupID int32) (*models.InvGroup, error)
	GetCategory(ctx context.Context, categoryID int32) (*models.InvCategory, error)
	GetMarketGroup(ctx context.Context, marketGroupID int32) (*models.InvMarketGroup, error)
	GetSolarSystem(ctx context.Context, systemID int32) (*models.MapSolarSystem, error)
	GetConstellation(ctx context.Context, constellationID int32) (*models.MapConstellation, error)
	GetRegion(ctx context.Context, regionID int32) (*models.MapRegion, error)
}

type DBWrapper struct {
	ch driver.Conn
	db *gorm.DB // For SQLite (EVE static data)
}

var db *DBWrapper

func GetDB() (DB, error) {
	if db != nil {
		return db, nil
	}

	sdb, err := GetDBSqlite()
	if err != nil {
		return nil, fmt.Errorf("failed to connect to SQLite: %w", err)
	}

	conn, err := GetDBClickhouse()
	if err != nil {
		return nil, fmt.Errorf("failed to connect to ClickHouse: %w", err)
	}

	db = &DBWrapper{
		ch: conn,
		db: sdb,
	}
	err = db.Init()
	return db, err
}

func GetDBSqlite() (*gorm.DB, error) {
	return gorm.Open(sqlite.Open("sqlite-latest.sqlite"), &gorm.Config{
		NamingStrategy: schema.NamingStrategy{
			NoLowerCase: true,
		},
	})
}

func GetDBClickhouse() (driver.Conn, error) {
	options := &clickhouse.Options{
		Addr: []string{"clickhouse.site.quack-lab.dev"},
		Auth: clickhouse.Auth{
			Database: "zkill",
			Username: "default",
			Password: "",
		},
		Protocol: clickhouse.HTTP,
		Settings: clickhouse.Settings{
			"max_query_size": 100000000,
		},
	}
	return clickhouse.Open(options)
}

func (db *DBWrapper) Get() *gorm.DB {
	return db.db
}

func (db *DBWrapper) Init() error {
	ctx := context.Background()

	// Migrate unified cache table
	// Use raw SQL to create table and index with IF NOT EXISTS to avoid errors
	// For 404s, we store a special marker byte sequence instead of NULL
	err := db.db.AutoMigrate(&CacheEntry{})
	if err != nil {
		return fmt.Errorf("failed to migrate cache_entries table: %w", err)
	}

	// Create killmails table
	createFlatKillmails := `
	CREATE TABLE IF NOT EXISTS killmails (
		killmail_id Int64,
		killmail_time DateTime,
		solar_system_id Int64,
		killmail_hash String,
		victim_ship_type_id Int64,
		victim_character_id Int64,
		victim_corporation_id Int64,
		victim_alliance_id Int64,
		victim_damage_taken Int64,
		victim_pos_x Float64,
		victim_pos_y Float64,
		victim_pos_z Float64,
		attacker_count UInt16,
		total_damage_done Int64,
		final_blow_ship_type Int64,
		attackers Array(Tuple(
			Int64,  -- character_id
			Int64,  -- corporation_id
			Int64,  -- alliance_id
			Int64,  -- ship_type_id
			Int64,  -- weapon_type_id
			Int64,  -- damage_done
			UInt8,  -- final_blow
			Float64 -- security_status
		)),
		items Array(Tuple(
			Int64,  -- flag
			Int64,  -- item_type_id
			Int64,  -- quantity_destroyed
			Int64,  -- quantity_dropped
			Int64   -- singleton
		))
	) ENGINE = MergeTree()
	ORDER BY (killmail_id)
	PRIMARY KEY (killmail_id)`
	err = db.ch.Exec(ctx, createFlatKillmails)
	if err != nil {
		return fmt.Errorf("failed to create killmails table: %w", err)
	}

	// Create modules table
	createFittedModules := `
	CREATE TABLE IF NOT EXISTS modules (
		killmail_id Int64,
		item_type_id Int64,
		slot String
	) ENGINE = MergeTree()
	ORDER BY (killmail_id, item_type_id)
	PRIMARY KEY (killmail_id, item_type_id)`
	err = db.ch.Exec(ctx, createFittedModules)
	if err != nil {
		return fmt.Errorf("failed to create modules table: %w", err)
	}

	return nil
}

func (db *DBWrapper) QueryFits(params QueryParams) (*FitStatistics, error) {
	ctx := context.Background()

	// Expand groups into item type IDs
	newItemTypes, err := db.ExpandGroupsIntoItemTypeIds(params.Groups)
	if err != nil {
		return nil, err
	}
	params.Modules = append(params.Modules, newItemTypes...)
	modules := deduplicateInt64(params.Modules)

	// Build different filter queries for different aggregations
	baseFilterQuery, baseFilterArgs := db.buildBaseFilterQuery(params)      // ship/system only
	fullFilterQuery, fullFilterArgs := db.buildFilterQuery(params, modules) // includes modules

	stats := &FitStatistics{
		ShipBreakdown:   make(map[int64]int64),
		SystemBreakdown: make(map[int64]int64),
		HighSlotModules: make(map[int32]int64),
		MidSlotModules:  make(map[int32]int64),
		LowSlotModules:  make(map[int32]int64),
		Rigs:            make(map[int32]int64),
		Drones:          make(map[int32]int64),
		KillmailIDs:     make([]int64, 0, params.KillmailLimit),
	}

	// Execute queries in parallel
	errChan := make(chan error, 5)

	// Query 1: Get total count (base filters only)
	go func() {
		query := fmt.Sprintf("SELECT COUNT(*) FROM (%s)", baseFilterQuery)
		rows, err := db.ch.Query(ctx, query, baseFilterArgs...)
		if err != nil {
			errChan <- fmt.Errorf("total count query failed: %w", err)
			return
		}
		defer rows.Close()

		if rows.Next() {
			var count uint64
			if err := rows.Scan(&count); err != nil {
				errChan <- fmt.Errorf("failed to scan total count: %w", err)
				return
			}
			stats.TotalKillmails = int64(count)
		}
		errChan <- nil
	}()

	// Query 2: Get killmail IDs (full filters)
	go func() {
		query := fmt.Sprintf("SELECT killmail_id FROM (%s) ORDER BY killmail_id LIMIT %d", fullFilterQuery, params.KillmailLimit)
		rows, err := db.ch.Query(ctx, query, fullFilterArgs...)
		if err != nil {
			errChan <- fmt.Errorf("killmail IDs query failed: %w", err)
			return
		}
		defer rows.Close()

		for rows.Next() {
			var id int64
			if err := rows.Scan(&id); err != nil {
				errChan <- fmt.Errorf("failed to scan killmail ID: %w", err)
				return
			}
			stats.KillmailIDs = append(stats.KillmailIDs, id)
		}
		errChan <- nil
	}()

	// Query 3: Ship breakdown (full filters)
	go func() {
		query := fmt.Sprintf("SELECT victim_ship_type_id, COUNT(*) FROM (%s) GROUP BY victim_ship_type_id", fullFilterQuery)
		rows, err := db.ch.Query(ctx, query, fullFilterArgs...)
		if err != nil {
			errChan <- fmt.Errorf("ship breakdown query failed: %w", err)
			return
		}
		defer rows.Close()

		for rows.Next() {
			var id int64
			var shipCount uint64
			if err := rows.Scan(&id, &shipCount); err != nil {
				errChan <- fmt.Errorf("failed to scan ship breakdown row: %w", err)
				return
			}
			stats.ShipBreakdown[id] = int64(shipCount)
		}
		errChan <- nil
	}()

	// Query 4: System breakdown (full filters)
	go func() {
		query := fmt.Sprintf("SELECT solar_system_id, COUNT(*) FROM (%s) GROUP BY solar_system_id", fullFilterQuery)
		rows, err := db.ch.Query(ctx, query, fullFilterArgs...)
		if err != nil {
			errChan <- fmt.Errorf("system breakdown query failed: %w", err)
			return
		}
		defer rows.Close()

		for rows.Next() {
			var id int64
			var count uint64
			if err := rows.Scan(&id, &count); err != nil {
				errChan <- fmt.Errorf("failed to scan system breakdown row: %w", err)
				return
			}
			stats.SystemBreakdown[id] = int64(count)
		}
		errChan <- nil
	}()

	// Query 5: Module statistics (full filters)
	go func() {
		subQuery := fmt.Sprintf("SELECT killmail_id FROM (%s)", fullFilterQuery)
		query := fmt.Sprintf(`
			SELECT slot, item_type_id, COUNT(DISTINCT killmail_id)
			FROM modules
			WHERE killmail_id IN (%s)
			GROUP BY slot, item_type_id`, subQuery)

		// Build args: fullFilterArgs + fullFilterArgs again for the subquery
		args := make([]interface{}, len(fullFilterArgs)*2)
		copy(args, fullFilterArgs)
		copy(args[len(fullFilterArgs):], fullFilterArgs)

		rows, err := db.ch.Query(ctx, query, args...)
		if err != nil {
			errChan <- fmt.Errorf("module stats query failed: %w", err)
			return
		}
		defer rows.Close()

		for rows.Next() {
			var slot string
			var itemID int64
			var count uint64

			if err := rows.Scan(&slot, &itemID, &count); err != nil {
				errChan <- fmt.Errorf("failed to scan module row: %w", err)
				return
			}
			int32ID := int32(itemID)
			int64Count := int64(count)

			switch slot {
			case "High":
				stats.HighSlotModules[int32ID] = int64Count
			case "Mid":
				stats.MidSlotModules[int32ID] = int64Count
			case "Low":
				stats.LowSlotModules[int32ID] = int64Count
			case "Rig":
				stats.Rigs[int32ID] = int64Count
			case "Drone":
				stats.Drones[int32ID] = int64Count
			}
		}
		errChan <- nil
	}()

	// Wait for all queries to complete
	for i := 0; i < 5; i++ {
		if err := <-errChan; err != nil {
			return nil, err
		}
	}

	// DISABLE CACHING TEMPORARILY TO TEST
	// // Cache the results
	// jsonData, err := json.Marshal(stats)
	// if err != nil {
	// 	flog.Debug("Failed to marshal results for caching: %v", err)
	// } else {
	// 	if cacheErr := db.CacheSet(cacheKey, jsonData); cacheErr != nil {
	// 		flog.Debug("Failed to cache results: %v", cacheErr)
	// 	} else {
	// 		flog.Debug("Cached results for key: %s", cacheKey)
	// 	}
	// }

	return stats, nil
}

func (db *DBWrapper) buildBaseFilterQuery(params QueryParams) (string, []interface{}) {
	query := `SELECT killmail_id, solar_system_id, victim_ship_type_id FROM killmails`
	args := []interface{}{}
	conditions := []string{}

	if params.Ship > 0 {
		conditions = append(conditions, "victim_ship_type_id = ?")
		args = append(args, params.Ship)
	}

	if len(params.Systems) > 0 {
		placeholders := make([]string, len(params.Systems))
		for i := range placeholders {
			placeholders[i] = "?"
		}
		conditions = append(conditions, "solar_system_id IN ("+strings.Join(placeholders, ",")+")")
		for _, sys := range params.Systems {
			args = append(args, sys)
		}
	}

	if len(conditions) > 0 {
		query += " WHERE " + strings.Join(conditions, " AND ")
	}

	return query, args
}

func (db *DBWrapper) buildFilterQuery(params QueryParams, modules []int64) (string, []interface{}) {
	query := `SELECT killmail_id, solar_system_id, victim_ship_type_id FROM zkill.killmails`
	args := []interface{}{}
	conditions := []string{}

	if len(modules) > 0 {
		query = `SELECT fk.killmail_id, fk.solar_system_id, fk.victim_ship_type_id FROM zkill.killmails fk INNER JOIN zkill.modules fm ON fk.killmail_id = fm.killmail_id`
		placeholders := make([]string, len(modules))
		for i := range placeholders {
			placeholders[i] = "?"
		}
		conditions = append(conditions, "fm.item_type_id IN ("+strings.Join(placeholders, ",")+")")
		for _, mod := range modules {
			args = append(args, mod)
		}
	}

	if params.Ship > 0 {
		conditions = append(conditions, "victim_ship_type_id = ?")
		args = append(args, params.Ship)
	}

	if len(params.Systems) > 0 {
		placeholders := make([]string, len(params.Systems))
		for i := range placeholders {
			placeholders[i] = "?"
		}
		conditions = append(conditions, "solar_system_id IN ("+strings.Join(placeholders, ",")+")")
		for _, sys := range params.Systems {
			args = append(args, sys)
		}
	}

	if len(conditions) > 0 {
		query += " WHERE " + strings.Join(conditions, " AND ")
	}

	return query, args
}

func (db *DBWrapper) ExpandGroupsIntoItemTypeIds(groups []int64) ([]int64, error) {
	var groupTypeIDs []int64
	result := db.db.Model(&models.InvType{}).
		Select("typeID").
		Where("groupID IN ?", groups).
		Pluck("typeID", &groupTypeIDs)
	return groupTypeIDs, result.Error
}

func (db *DBWrapper) SearchShips(query string, limit int) ([]models.InvType, error) {
	var ships []models.InvType
	searchPattern := "%" + strings.ToLower(query) + "%"
	err := db.db.Table("invTypes").
		Joins("INNER JOIN invGroups ON invTypes.groupID = invGroups.groupID").
		Where("LOWER(invTypes.\"typeName\") LIKE ? AND invGroups.categoryID IN (6)", searchPattern).
		Limit(limit).
		Find(&ships).Error
	return ships, err
}

func (db *DBWrapper) SearchSystems(query string, limit int) ([]models.MapSolarSystem, error) {
	var systems []models.MapSolarSystem
	searchPattern := "%" + strings.ToLower(query) + "%"
	err := db.db.Table("mapSolarSystems").
		Where("LOWER(\"solarSystemName\") LIKE ?", searchPattern).
		Limit(limit).
		Find(&systems).Error
	return systems, err
}

func (db *DBWrapper) SearchModules(query string, limit int) ([]models.InvType, error) {
	var modules []models.InvType
	searchPattern := "%" + strings.ToLower(query) + "%"
	err := db.db.Table("invTypes").
		Joins("INNER JOIN invGroups ON invTypes.groupID = invGroups.groupID").
		Where("LOWER(invTypes.\"typeName\") LIKE ? AND invGroups.categoryID IN (7, 66)", searchPattern).
		Limit(limit).
		Find(&modules).Error
	return modules, err
}

func (db *DBWrapper) SearchGroups(query string, limit int) ([]models.InvGroup, error) {
	var groups []models.InvGroup
	searchPattern := "%" + strings.ToLower(query) + "%"
	err := db.db.Table("invGroups").
		Where("LOWER(\"groupName\") LIKE ?", searchPattern).
		Limit(limit).
		Find(&groups).Error
	return groups, err
}

func (db *DBWrapper) GetItemTypes(itemIDs []int64) ([]models.InvType, error) {
	var itemTypes []models.InvType
	res := db.db.Model(&models.InvType{}).
		Where("typeID IN ?", itemIDs).
		Find(&itemTypes)
	return itemTypes, res.Error
}

func (db *DBWrapper) GetSolarSystems(systemIDs []int64) ([]models.MapSolarSystem, error) {
	var systems []models.MapSolarSystem
	res := db.db.Model(&models.MapSolarSystem{}).
		Where("solarSystemID IN ?", systemIDs).
		Find(&systems)
	return systems, res.Error
}

func deduplicateInt64(slice []int64) []int64 {
	seen := make(map[int64]bool)
	result := make([]int64, 0, len(slice))
	for _, v := range slice {
		if !seen[v] {
			seen[v] = true
			result = append(result, v)
		}
	}
	return result
}

func (db *DBWrapper) CacheSet(key string, data []byte) error {
	flog := logger.Default.WithPrefix("CacheSet").WithPrefix(key)
	flog.Dump("data", data)

	err := db.CacheClean()
	if err != nil {
		flog.Debug("Failed to clean cache: %v", err)
		return err
	}

	cacheEntry := CacheEntry{
		Key:       key,
		Data:      data,
		CreatedAt: time.Now(),
	}
	flog.Debug("Creating cache entry")
	flog.Dump("cacheEntry", cacheEntry)
	return db.db.Create(&cacheEntry).Error
}

var ErrCacheMiss = gorm.ErrRecordNotFound

func (db *DBWrapper) CacheGet(key string) ([]byte, error) {
	flog := logger.Default.WithPrefix("CacheGet").WithPrefix(key)
	flog.Dump("key", key)

	err := db.CacheClean()
	if err != nil {
		flog.Debug("Failed to clean cache: %v", err)
		return nil, err
	}

	cacheEntry := CacheEntry{Key: key}
	res := db.db.Model(&CacheEntry{}).
		Where(cacheEntry).
		First(&cacheEntry)
	flog.Debug("Found cache entry")
	flog.Dump("cacheEntry", cacheEntry)
	flog.Dump("res", res)
	return cacheEntry.Data, res.Error
}

func (db *DBWrapper) CacheClean() error {
	flog := logger.Default.WithPrefix("CacheClean")
	threshold := time.Now().Add(-72 * time.Hour)
	flog.Dump("threshold", threshold)
	return db.db.
		Where("created_at < ?", threshold).
		Delete(&CacheEntry{}).Error
}

func (db *DBWrapper) GetModuleSlots(moduleIDs []int64) (map[int64]ModuleSlot, error) {
	if len(moduleIDs) == 0 {
		return make(map[int64]ModuleSlot), nil
	}

	var effects []models.DgmTypeEffect
	qres := db.db.Model(&models.DgmTypeEffect{}).
		Select("typeID, effectID").
		Where("typeID IN ? AND effectID IN (11, 12, 13, 2663)", moduleIDs).
		Find(&effects)
	if qres.Error != nil {
		return nil, qres.Error
	}

	result := make(map[int64]ModuleSlot)
	for _, e := range effects {
		var slot ModuleSlot
		switch e.EffectID {
		case 11:
			slot = ModuleSlotLow
		case 12:
			slot = ModuleSlotHigh
		case 13:
			slot = ModuleSlotMid
		case 2663:
			slot = ModuleSlotRig
		}
		result[int64(e.TypeID)] = slot
	}

	return result, nil
}

func limitKillmails(killmailIDs []int64, limit int) []int64 {
	if limit <= 0 || len(killmailIDs) <= limit {
		return killmailIDs
	}
	return killmailIDs[:limit]
}

// parseKeyValuePairs parses a string like "123:456,789:012" into key-value pairs
func parseKeyValuePairs(data string, callback func(key, value int64)) {
	if data == "" {
		return
	}

	pairs := strings.Split(data, ",")
	for _, pair := range pairs {
		if kv := strings.Split(pair, ":"); len(kv) == 2 {
			if key, err := strconv.ParseInt(kv[0], 10, 64); err == nil {
				if value, err := strconv.ParseInt(kv[1], 10, 64); err == nil {
					callback(key, value)
				}
			}
		}
	}
}

// parseKillmailIDs parses a string like "123,456,789" into a slice of int64
func parseKillmailIDs(data string, result *[]int64) {
	if data == "" {
		return
	}

	ids := strings.Split(data, ",")
	for _, idStr := range ids {
		if id, err := strconv.ParseInt(idStr, 10, 64); err == nil {
			*result = append(*result, id)
		}
	}
}

// HasModules returns true if the query has module filters
func (qp QueryParams) HasModules() bool {
	return len(qp.Modules) > 0
}

func (db *DBWrapper) GetType(ctx context.Context, typeID int32) (*models.InvType, error) {
	var t models.InvType
	if err := db.db.Where("typeID = ?", typeID).First(&t).Error; err != nil {
		return nil, fmt.Errorf("failed to get type %d: %w", typeID, err)
	}
	return &t, nil
}

func (db *DBWrapper) GetGroup(ctx context.Context, groupID int32) (*models.InvGroup, error) {
	var g models.InvGroup
	if err := db.db.Where("groupID = ?", groupID).First(&g).Error; err != nil {
		return nil, fmt.Errorf("failed to get group %d: %w", groupID, err)
	}
	return &g, nil
}

func (db *DBWrapper) GetCategory(ctx context.Context, categoryID int32) (*models.InvCategory, error) {
	var c models.InvCategory
	if err := db.db.Where("categoryID = ?", categoryID).First(&c).Error; err != nil {
		return nil, fmt.Errorf("failed to get category %d: %w", categoryID, err)
	}
	return &c, nil
}

func (db *DBWrapper) GetMarketGroup(ctx context.Context, marketGroupID int32) (*models.InvMarketGroup, error) {
	var mg models.InvMarketGroup
	if err := db.db.Where("marketGroupID = ?", marketGroupID).First(&mg).Error; err != nil {
		return nil, fmt.Errorf("failed to get market group %d: %w", marketGroupID, err)
	}
	return &mg, nil
}

func (db *DBWrapper) GetSolarSystem(ctx context.Context, systemID int32) (*models.MapSolarSystem, error) {
	var s models.MapSolarSystem
	if err := db.db.Where("solarSystemID = ?", systemID).First(&s).Error; err != nil {
		return nil, fmt.Errorf("failed to get solar system %d: %w", systemID, err)
	}
	return &s, nil
}

func (db *DBWrapper) GetConstellation(ctx context.Context, constellationID int32) (*models.MapConstellation, error) {
	var c models.MapConstellation
	if err := db.db.Where("constellationID = ?", constellationID).First(&c).Error; err != nil {
		return nil, fmt.Errorf("failed to get constellation %d: %w", constellationID, err)
	}
	return &c, nil
}

func (db *DBWrapper) GetRegion(ctx context.Context, regionID int32) (*models.MapRegion, error) {
	var r models.MapRegion
	if err := db.db.Where("regionID = ?", regionID).First(&r).Error; err != nil {
		return nil, fmt.Errorf("failed to get region %d: %w", regionID, err)
	}
	return &r, nil
}