package main

import (
	"archive/tar"
	"context"
	"database/sql"
	"io"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

func TestRunningDaemonBackupRestoresTheChatAndLinks(t *testing.T) {
	dir := t.TempDir()
	whatsappTestSession(t, dir, true)

	files := map[string]string{
		"auth.json":                  `{"access_token":"fake-secret"}`,
		"telegram.json":              `{"token":"fake-token","owner":123}`,
		"config.json":                `{"endpoint":"api"}`,
		"agent/instructions.md":      "Keep the instructions.\n",
		"agent/telegram-offset.json": "23\n",
	}
	for name, data := range files {
		path := filepath.Join(dir, name)
		require.NoError(t, os.MkdirAll(filepath.Dir(path), 0o700))
		require.NoError(t, os.WriteFile(path, []byte(data), 0o600))
	}

	memory, whatsapp, subagents := backupStartDaemon(t, dir)
	require.True(t, whatsapp.isLinked())

	_, err := memory.Append("note", "A durable decision.")
	require.NoError(t, err)

	_, err = memory.Append("note", "The next decision.")
	require.NoError(t, err)
	require.NoError(t, memory.Settle(t.Context()))

	child, err := subagents.Create(t.Context())
	require.NoError(t, err)

	_, err = child.Append("note", "subagent")
	require.NoError(t, err)

	_, err = child.Append("user", "A saved task.")
	require.NoError(t, err)

	childText, err := child.Show()
	require.NoError(t, err)

	for _, excluded := range []string{"probe/skip", "agent-backup-example/skip"} {
		path := filepath.Join(dir, excluded)
		require.NoError(t, os.MkdirAll(filepath.Dir(path), 0o700))
		require.NoError(t, os.WriteFile(path, []byte("excluded"), 0o600))
	}

	archive := filepath.Join(t.TempDir(), "daily.tar.gz")
	t.Setenv("ONECHAT_DIR", dir)

	var answer strings.Builder

	require.NoError(t, backup(t.Context(), archive, &answer))
	require.Regexp(t, `^`+archive+` \([0-9]+ bytes\)\n$`, answer.String())

	info, err := os.Stat(archive)
	require.NoError(t, err)
	require.Equal(t, os.FileMode(0o600), info.Mode().Perm())

	restored := t.TempDir()
	require.NoError(t, backupRestore(t.Context(), archive, restored))

	for name, data := range files {
		bytes, err := os.ReadFile(filepath.Join(restored, name))
		require.NoError(t, err)
		require.Equal(t, data, string(bytes))
	}

	for _, excluded := range []string{"probe", "agent-backup-example", socketFileName, "agent/lock"} {
		_, err = os.Stat(filepath.Join(restored, excluded))
		require.ErrorIs(t, err, os.ErrNotExist)
	}

	restoredMemory, restoredWhatsApp, restoredSubagents := backupStartDaemon(t, restored)
	require.Equal(t, memory.Len(), restoredMemory.Len())
	require.Equal(t, memory.View(), restoredMemory.View())

	for id := range memory.Len() {
		expected, err := memory.Message(id)
		require.NoError(t, err)

		actual, err := restoredMemory.Message(id)
		require.NoError(t, err)
		require.Equal(t, expected, actual)
	}

	require.True(t, restoredWhatsApp.isLinked())
	require.Equal(t, whatsapp.link().own, restoredWhatsApp.link().own)

	restoredChild, err := restoredSubagents.Get(t.Context(), child.name)
	require.NoError(t, err)

	restoredText, err := restoredChild.Show()
	require.NoError(t, err)
	require.Equal(t, childText, restoredText)
}

func TestBackupRestoresEachChildWithItsOwnState(t *testing.T) {
	dir := t.TempDir()
	daemonChildConfig(t, dir, "Garden")
	daemonChildConfig(t, dir, "Library")

	server := newModelFake(t)
	clock := newTestClock()
	agents, err := openDaemonAgents(t.Context(), dir, server.URL, ModelAPIKey(modelTestKey), defaultConfig().Compactor, clock.Now, discardLog())
	require.NoError(t, err)

	defer func() {
		require.NoError(t, agents.Close())
	}()

	whatsapp, _, _ := whatsappStartTest(t, dir, clock)

	var loops []*TurnLoop

	for _, agent := range agents {
		loop, err := agent.Loop(whatsapp, func(Event) {
		})
		require.NoError(t, err)

		loops = append(loops, loop)
		_, err = loop.Submit("parent", "Keep the decision of "+agent.agent.config.Name)
		require.NoError(t, err)
		require.NoError(t, loop.turn(t.Context()))
		require.NoError(t, agent.memory.Settle(t.Context()))

		subagent, err := agent.subagents.Create(t.Context())
		require.NoError(t, err)

		_, err = subagent.Append("user", "Keep the task of "+agent.agent.config.Name)
		require.NoError(t, err)
	}

	archive := filepath.Join(t.TempDir(), "children.tar.gz")
	_, err = backupData(t.Context(), dir, archive, whatsapp, func(ctx context.Context, snapshot func() error) error {
		return daemonSnapshot(ctx, loops, snapshot)
	})
	require.NoError(t, err)

	restored := t.TempDir()
	require.NoError(t, backupRestore(t.Context(), archive, restored))

	restoredAgents, err := openDaemonAgents(t.Context(), restored, server.URL, ModelAPIKey(modelTestKey), defaultConfig().Compactor, clock.Now, discardLog())
	require.NoError(t, err)

	defer func() {
		require.NoError(t, restoredAgents.Close())
	}()

	require.Len(t, restoredAgents, 3)

	for i, agent := range agents {
		other := restoredAgents[i]
		require.Equal(t, agent.agent.config, other.agent.config)
		require.Equal(t, agent.memory.View(), other.memory.View())
		require.Equal(t, agent.memory.Len(), other.memory.Len())
		require.Equal(t, modelReadUsage(t, agent.client.UsagePath()), modelReadUsage(t, other.client.UsagePath()))

		for id := range agent.memory.Len() {
			before, err := agent.memory.Message(id)
			require.NoError(t, err)

			after, err := other.memory.Message(id)
			require.NoError(t, err)
			require.Equal(t, before, after)
		}

		entries, err := os.ReadDir(filepath.Join(agent.agent.dir, subagentDirName))
		require.NoError(t, err)
		require.Len(t, entries, 1)

		subagent, err := other.subagents.Get(t.Context(), entries[0].Name())
		require.NoError(t, err)

		text, err := subagent.Show()
		require.NoError(t, err)
		require.Contains(t, text, "Keep the task of "+agent.agent.config.Name)
	}
}

func TestRestoreRefusesNonemptyDataWithoutChangingIt(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "keep")
	require.NoError(t, os.WriteFile(path, []byte("keep"), 0o600))

	err := backupRestore(t.Context(), "missing.tar.gz", dir)
	require.ErrorContains(t, err, "empty data directory")

	entries, err := os.ReadDir(dir)
	require.NoError(t, err)
	require.Len(t, entries, 1)

	data, err := os.ReadFile(path)
	require.NoError(t, err)
	require.Equal(t, "keep", string(data))
}

func TestRestoreRefusesTheDaemonTokenLock(t *testing.T) {
	dir := t.TempDir()
	tokens, err := lockTokens(dir)
	require.NoError(t, err)

	defer func() {
		require.NoError(t, tokens.Close())
	}()

	err = backupRestore(t.Context(), "missing.tar.gz", dir)
	require.ErrorIs(t, err, errTokensInUse)
}

func TestUnsafeRestoreInstallsNoArchiveFiles(t *testing.T) {
	archive := backupTestArchive(t,
		tar.Header{
			Name:     "safe",
			Typeflag: tar.TypeReg,
			Mode:     0o600,
		},
		tar.Header{
			Name:     "../escape",
			Typeflag: tar.TypeReg,
			Mode:     0o600,
		},
	)
	dir := t.TempDir()
	err := backupRestore(t.Context(), archive, dir)
	require.ErrorContains(t, err, "unsafe archive path")

	entries, err := os.ReadDir(dir)
	require.NoError(t, err)
	require.Len(t, entries, 1)
	require.Equal(t, tokenLockFileName, entries[0].Name())
}

func TestBackupNeedsARunningDaemon(t *testing.T) {
	t.Setenv("ONECHAT_DIR", t.TempDir())

	err := backup(t.Context(), filepath.Join(t.TempDir(), "backup.tar.gz"), io.Discard)
	require.EqualError(t, err, "backup needs a running daemon")
}

func TestWhatsAppSessionBackupIncludesCommittedWALWrites(t *testing.T) {
	dir := t.TempDir()
	whatsappTestSession(t, dir, true)

	whatsapp, _, _ := whatsappStartTest(t, dir, newTestClock())
	address, err := whatsappSessionAddress(filepath.Join(dir, whatsappDirName, whatsappSessionFileName))
	require.NoError(t, err)

	database, err := sql.Open("sqlite", address)
	require.NoError(t, err)

	defer func() {
		require.NoError(t, database.Close())
	}()

	_, err = database.ExecContext(t.Context(), "PRAGMA journal_mode=WAL; CREATE TABLE backup_sample (value TEXT); INSERT INTO backup_sample VALUES ('saved');")
	require.NoError(t, err)

	target := filepath.Join(t.TempDir(), "copy", whatsappSessionFileName)
	require.NoError(t, whatsapp.BackupSession(t.Context(), target))

	address, err = whatsappSessionAddress(target)
	require.NoError(t, err)

	restored, err := sql.Open("sqlite", address)
	require.NoError(t, err)

	defer func() {
		require.NoError(t, restored.Close())
	}()

	var value string

	err = restored.QueryRowContext(t.Context(), "SELECT value FROM backup_sample").Scan(&value)
	require.NoError(t, err)
	require.Equal(t, "saved", value)

	info, err := os.Stat(target)
	require.NoError(t, err)
	require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
}

func TestSessionBackupWaitsForRemovalToFinish(t *testing.T) {
	dir := t.TempDir()
	whatsappTestSession(t, dir, true)

	whatsapp, fake, _ := whatsappStartTest(t, dir, newTestClock())
	disconnect := &BackupDisconnect{
		WhatsAppClient: fake,
		entered:        make(chan struct{}),
		release:        make(chan struct{}),
	}

	var release sync.Once

	defer release.Do(func() {
		close(disconnect.release)
	})

	whatsapp.mutex.Lock()

	whatsapp.client = disconnect
	whatsapp.mutex.Unlock()

	removed := make(chan error, 1)
	go func() {
		removed <- whatsapp.removeSession()
	}()

	<-disconnect.entered

	target := filepath.Join(t.TempDir(), "copy", whatsappSessionFileName)
	backedUp := make(chan error, 1)
	started := make(chan struct{})
	go func() {
		close(started)

		backedUp <- whatsapp.BackupSession(t.Context(), target)
	}()

	<-started
	require.Never(t, func() bool {
		return len(backedUp) != 0
	}, 50*time.Millisecond, time.Millisecond)
	release.Do(func() {
		close(disconnect.release)
	})
	require.NoError(t, <-removed)
	require.NoError(t, <-backedUp)

	for _, path := range []string{whatsapp.sessionPath(), target} {
		_, err := os.Stat(path)
		require.ErrorIs(t, err, os.ErrNotExist)
	}
}

// BackupDisconnect holds session removal in its unlocked client cleanup.
type BackupDisconnect struct {
	WhatsAppClient
	entered chan struct{}
	release chan struct{}
}

func (d *BackupDisconnect) Disconnect() {
	close(d.entered)
	<-d.release
}

// backupStartDaemon serves fake model and WhatsApp clients on the socket of
// dir. It keeps all state files under dir and holds the daemon token lock.
func backupStartDaemon(t *testing.T, dir string) (*Memory, *WhatsApp, *Subagents) {
	t.Helper()

	tokens, err := lockTokens(dir)
	require.NoError(t, err)

	clock := newTestClock()
	agent, err := openAgent(filepath.Join(dir, daemonAgentDirName))
	require.NoError(t, err)

	memory, err := openMemory(t.Context(), agent.dir, clock.Now, newFakeSummarizer(), discardLog())
	require.NoError(t, err)

	whatsapp, _, _ := whatsappStartTest(t, dir, clock)
	subagents := newSubagents(filepath.Join(agent.dir, subagentDirName), clock.Now, subagentNameRandom)
	schedules, err := openSchedules(filepath.Join(agent.dir, scheduleFileName), "OneChat", clock.Now, discardLog())
	require.NoError(t, err)

	server := newModelFake(t)
	client := newModelClient(server.URL, ModelAPIKey(modelTestKey), agent.CacheKey(), filepath.Join(agent.dir, "usage.jsonl"))
	listener, err := daemonListenSocket(filepath.Join(dir, socketFileName))
	require.NoError(t, err)

	ctx, cancel := context.WithCancel(t.Context())
	done := make(chan error, 1)
	go func() {
		done <- daemonServe(ctx, listener, nil, whatsapp, schedules, func(notify func(Event)) (*TurnLoop, error) {
			return newTurnLoop(agent, memory, subagents, whatsapp, schedules, client, clock.Now, discardLog(), notify)
		})
	}()

	t.Cleanup(func() {
		cancel()
		require.NoError(t, <-done)
		require.NoError(t, subagents.Close())
		require.NoError(t, memory.Close())
		require.NoError(t, tokens.Close())
	})

	return memory, whatsapp, subagents
}