package storage

import (
	"context"
	"sync"
	"time"

	"github.com/google/uuid"
)

// MemoryStore is an in-memory implementation of Store for use in tests.
type MemoryStore struct {
	mu            sync.Mutex
	LoginAttempts []LoginAttempt
	Sessions      map[string]*Session
	SessionLogs   []SessionLog
}

// NewMemoryStore returns a new empty MemoryStore.
func NewMemoryStore() *MemoryStore {
	return &MemoryStore{
		Sessions: make(map[string]*Session),
	}
}

func (m *MemoryStore) RecordLoginAttempt(_ context.Context, username, password, ip string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	now := time.Now().UTC()
	for i := range m.LoginAttempts {
		a := &m.LoginAttempts[i]
		if a.Username == username && a.Password == password && a.IP == ip {
			a.Count++
			a.LastSeen = now
			return nil
		}
	}

	m.LoginAttempts = append(m.LoginAttempts, LoginAttempt{
		ID:        int64(len(m.LoginAttempts) + 1),
		Username:  username,
		Password:  password,
		IP:        ip,
		Count:     1,
		FirstSeen: now,
		LastSeen:  now,
	})
	return nil
}

func (m *MemoryStore) CreateSession(_ context.Context, ip, username, shellName string) (string, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	id := uuid.New().String()
	now := time.Now().UTC()
	m.Sessions[id] = &Session{
		ID:          id,
		IP:          ip,
		Username:    username,
		ShellName:   shellName,
		ConnectedAt: now,
	}
	return id, nil
}

func (m *MemoryStore) EndSession(_ context.Context, sessionID string, disconnectedAt time.Time) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if s, ok := m.Sessions[sessionID]; ok {
		t := disconnectedAt.UTC()
		s.DisconnectedAt = &t
	}
	return nil
}

func (m *MemoryStore) UpdateHumanScore(_ context.Context, sessionID string, score float64) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	if s, ok := m.Sessions[sessionID]; ok {
		s.HumanScore = &score
	}
	return nil
}

func (m *MemoryStore) AppendSessionLog(_ context.Context, sessionID, input, output string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	m.SessionLogs = append(m.SessionLogs, SessionLog{
		ID:        int64(len(m.SessionLogs) + 1),
		SessionID: sessionID,
		Timestamp: time.Now().UTC(),
		Input:     input,
		Output:    output,
	})
	return nil
}

func (m *MemoryStore) DeleteRecordsBefore(_ context.Context, cutoff time.Time) (int64, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	var total int64

	// Delete old login attempts.
	kept := m.LoginAttempts[:0]
	for _, a := range m.LoginAttempts {
		if a.LastSeen.Before(cutoff) {
			total++
		} else {
			kept = append(kept, a)
		}
	}
	m.LoginAttempts = kept

	// Delete old sessions and their logs.
	for id, s := range m.Sessions {
		if s.ConnectedAt.Before(cutoff) {
			delete(m.Sessions, id)
			total++
		}
	}

	keptLogs := m.SessionLogs[:0]
	for _, l := range m.SessionLogs {
		if _, ok := m.Sessions[l.SessionID]; ok {
			keptLogs = append(keptLogs, l)
		} else {
			total++
		}
	}
	m.SessionLogs = keptLogs

	return total, nil
}

func (m *MemoryStore) Close() error {
	return nil
}