package store

import (
	"database/sql"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"time"
)

const (
	auditDayStride   int64 = 1 << 40
	auditShardStride int64 = 1 << 28
	auditMaxParts          = int(auditDayStride/auditShardStride) - 2
)
const auditSchema = `PRAGMA journal_mode=WAL; PRAGMA busy_timeout=5000; PRAGMA foreign_keys=ON;
CREATE TABLE IF NOT EXISTS commands(id TEXT PRIMARY KEY,proxy_id TEXT NOT NULL,target_id TEXT NOT NULL,source TEXT NOT NULL,protocol TEXT NOT NULL,kind TEXT NOT NULL,decision TEXT NOT NULL,content BLOB NOT NULL,at TEXT NOT NULL);
CREATE TABLE IF NOT EXISTS command_results(command_id TEXT PRIMARY KEY REFERENCES commands(id) ON DELETE CASCADE,content BLOB NOT NULL);
CREATE TABLE IF NOT EXISTS command_outputs(seq INTEGER PRIMARY KEY AUTOINCREMENT,id TEXT NOT NULL UNIQUE,command_id TEXT NOT NULL REFERENCES commands(id) ON DELETE CASCADE,stream TEXT NOT NULL,content BLOB NOT NULL,size INTEGER NOT NULL);
CREATE INDEX IF NOT EXISTS command_outputs_command ON command_outputs(command_id,seq);
CREATE INDEX IF NOT EXISTS commands_at ON commands(at);
CREATE INDEX IF NOT EXISTS commands_proxy ON commands(proxy_id);
CREATE INDEX IF NOT EXISTS commands_target ON commands(target_id);
CREATE INDEX IF NOT EXISTS commands_filter ON commands(protocol,decision,kind);
CREATE TABLE IF NOT EXISTS events(id INTEGER PRIMARY KEY AUTOINCREMENT,proxy_id TEXT NOT NULL,source TEXT NOT NULL,category TEXT NOT NULL,detail TEXT NOT NULL,at TEXT NOT NULL);
CREATE INDEX IF NOT EXISTS events_at ON events(at);
CREATE INDEX IF NOT EXISTS events_proxy ON events(proxy_id);`

func auditShardDate(key string) (time.Time, error) {
	base := key
	if index := strings.Index(base, ".part"); index >= 0 {
		part := base[index+len(".part"):]
		if part == "" {
			return time.Time{}, errors.New("invalid audit shard part")
		}
		value, err := strconv.Atoi(part)
		if err != nil {
			return time.Time{}, err
		}
		if value < 1 || value > auditMaxParts {
			return time.Time{}, errors.New("audit shard part is out of range")
		}
		base = base[:index]
	}
	return time.Parse("2006-01-02", base)
}

func auditShardPath(directory, key string) string {
	return filepath.Join(directory, key+".sqlite")
}

func auditShardPart(key string) int {
	if index := strings.Index(key, ".part"); index >= 0 {
		part, _ := strconv.Atoi(key[index+len(".part"):])
		return part
	}
	return 0
}

func auditShardBases(keys []string) map[string]int64 {
	maxPart := map[string]int{}
	for _, key := range keys {
		date, err := auditShardDate(key)
		if err != nil {
			continue
		}
		day := date.Format("2006-01-02")
		if part := auditShardPart(key); part > maxPart[day] {
			maxPart[day] = part
		}
	}
	bases := make(map[string]int64, len(keys))
	for _, key := range keys {
		date, err := auditShardDate(key)
		if err != nil {
			continue
		}
		part := auditShardPart(key)
		if part == 0 {
			part = maxPart[date.Format("2006-01-02")] + 1
		}
		bases[key] = date.Unix()/86400*auditDayStride + int64(part)*auditShardStride
	}
	return bases
}

func (s *Store) auditDB(day string) (*sql.DB, error) {
	if _, err := auditShardDate(day); err != nil {
		return nil, err
	}
	s.auditMu.Lock()
	defer s.auditMu.Unlock()
	if s.closed {
		return nil, errors.New("store is closed")
	}
	return s.openAuditDBLocked(day)
}

func (s *Store) openAuditDBLocked(day string) (*sql.DB, error) {
	if db := s.audits[day]; db != nil {
		return db, nil
	}
	if err := os.MkdirAll(s.auditDir, 0700); err != nil {
		return nil, err
	}
	db, err := sql.Open("sqlite", auditShardPath(s.auditDir, day))
	if err != nil {
		return nil, err
	}
	db.SetMaxOpenConns(1)
	if _, err = db.Exec(auditSchema); err != nil {
		db.Close()
		return nil, err
	}
	s.audits[day] = db
	return db, nil
}

// auditDBForWrite rotates the active day shard before it grows beyond the
// configured limit. A single record may exceed the limit; it is never split
// across databases, so command/output writes remain atomic.
func (s *Store) auditDBForWrite(day string, estimate int64) (*sql.DB, error) {
	if _, err := auditShardDate(day); err != nil {
		return nil, err
	}
	s.auditMu.Lock()
	defer s.auditMu.Unlock()
	if s.closed {
		return nil, errors.New("store is closed")
	}
	db, err := s.openAuditDBLocked(day)
	if err != nil {
		return nil, err
	}
	if estimate <= 0 || s.auditMaxBytes <= 0 {
		return db, nil
	}
	current, err := s.auditFileBytesLocked(day)
	if err != nil {
		return nil, err
	}
	if current == 0 || current+estimate <= s.auditMaxBytes {
		return db, nil
	}
	if err = s.rotateAuditLocked(day); err != nil {
		return nil, err
	}
	return s.openAuditDBLocked(day)
}

func (s *Store) auditFileBytesLocked(day string) (int64, error) {
	var total int64
	for _, suffix := range []string{"", "-wal", "-shm"} {
		info, err := os.Stat(auditShardPath(s.auditDir, day) + suffix)
		if errors.Is(err, os.ErrNotExist) {
			continue
		}
		if err != nil {
			return 0, err
		}
		total += info.Size()
	}
	return total, nil
}

func (s *Store) rotateAuditLocked(day string) error {
	db := s.audits[day]
	if db != nil {
		if _, err := db.Exec("PRAGMA wal_checkpoint(TRUNCATE)"); err != nil {
			return err
		}
		var journalMode string
		if err := db.QueryRow("PRAGMA journal_mode=DELETE").Scan(&journalMode); err != nil {
			return err
		}
		if !strings.EqualFold(journalMode, "delete") {
			return fmt.Errorf("close audit WAL before rotation: unexpected journal mode %q", journalMode)
		}
		if err := db.Close(); err != nil {
			return err
		}
		delete(s.audits, day)
	}
	path := auditShardPath(s.auditDir, day)
	if _, err := os.Stat(path); errors.Is(err, os.ErrNotExist) {
		return nil
	} else if err != nil {
		return err
	}
	// DELETE mode leaves the main file self-contained. Remove any stale WAL
	// sidecars before renaming so a crash cannot pair them with a new shard.
	for _, suffix := range []string{"-wal", "-shm"} {
		if err := os.Remove(path + suffix); err != nil && !errors.Is(err, os.ErrNotExist) {
			return err
		}
	}
	part, err := s.nextAuditPartLocked(day)
	if err != nil {
		return err
	}
	archived := auditShardPath(s.auditDir, day+".part"+formatAuditPart(part))
	if err := os.Rename(path, archived); err != nil {
		return err
	}
	return nil
}

func (s *Store) nextAuditPartLocked(day string) (int, error) {
	entries, err := os.ReadDir(s.auditDir)
	if err != nil {
		return 0, err
	}
	prefix := day + ".part"
	maxPart := 0
	for _, entry := range entries {
		if entry.IsDir() || !strings.HasPrefix(entry.Name(), prefix) || !strings.HasSuffix(entry.Name(), ".sqlite") {
			continue
		}
		key := strings.TrimSuffix(entry.Name(), ".sqlite")
		if _, err = auditShardDate(key); err != nil {
			continue
		}
		if part := auditShardPart(key); part > maxPart {
			maxPart = part
		}
	}
	if maxPart >= auditMaxParts {
		return 0, errors.New("daily audit shard limit exceeded")
	}
	return maxPart + 1, nil
}

func formatAuditPart(part int) string {
	return fmt.Sprintf("%04d", part)
}

func (s *Store) auditDays() ([]string, error) {
	entries, err := os.ReadDir(s.auditDir)
	if errors.Is(err, os.ErrNotExist) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	days := []string{}
	for _, entry := range entries {
		if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".sqlite") {
			continue
		}
		day := strings.TrimSuffix(entry.Name(), ".sqlite")
		if _, err = auditShardDate(day); err == nil {
			days = append(days, day)
		}
	}
	sort.Slice(days, func(i, j int) bool {
		left, _ := auditShardDate(days[i])
		right, _ := auditShardDate(days[j])
		if left.Equal(right) {
			leftPart, rightPart := auditShardPart(days[i]), auditShardPart(days[j])
			if leftPart == 0 {
				return rightPart != 0
			}
			if rightPart == 0 {
				return false
			}
			return leftPart > rightPart
		}
		return left.After(right)
	})
	return days, nil
}
func (s *Store) pruneCommands(cutoff time.Time) {
	s.auditOps.Lock()
	defer s.auditOps.Unlock()
	if s.closed {
		return
	}
	cutoff = cutoff.UTC()
	days, err := s.auditDays()
	if err != nil {
		return
	}
	for _, day := range days {
		date, parseErr := auditShardDate(day)
		if parseErr != nil {
			continue
		}
		dayDate := date.Format("2006-01-02")
		if dayDate < cutoff.Format("2006-01-02") {
			s.auditMu.Lock()
			if db := s.audits[day]; db != nil {
				db.Close()
				delete(s.audits, day)
			}
			for _, suffix := range []string{"", "-wal", "-shm"} {
				_ = os.Remove(auditShardPath(s.auditDir, day) + suffix)
			}
			s.auditMu.Unlock()
		} else if dayDate == cutoff.Format("2006-01-02") {
			db, err := s.auditDB(day)
			if err == nil {
				db.Exec("DELETE FROM commands WHERE julianday(at)<julianday(?)", cutoff.Format(time.RFC3339Nano))
				db.Exec("DELETE FROM events WHERE julianday(at)<julianday(?)", cutoff.Format(time.RFC3339Nano))
			}
		}
	}
}