package agents

import (
	"context"
	"log/slog"

	"g.fc.run/misc/byre/internal/api"
	"g.fc.run/misc/byre/internal/git"
)

type mergeTurn struct {
	dir    string
	unlock func()
}

func (m *Manager) lockTask(taskID string) func() {
	return m.lockDelivery("critical:" + taskID)
}

func (m *Manager) ownsTask(sessionID string) bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	_, found := m.mergeTurns[sessionID]
	return found
}

func (m *Manager) takeMergeTurn(l *live) {
	m.ownMerge(l.id, l.taskID, l.dir)
}

func (m *Manager) ownMerge(sessionID, taskID, dir string) {
	if m.ownsTask(sessionID) {
		return
	}
	unlock := m.lockTask(taskID)
	m.mu.Lock()
	m.mergeTurns[sessionID] = mergeTurn{dir: dir, unlock: unlock}
	m.mu.Unlock()
}

func (m *Manager) recoverMergeTurns(ctx context.Context, running []api.AgentSession) error {
	for _, sess := range running {
		if sess.Role != api.Implementer {
			continue
		}
		lead, found, err := m.liveSession(ctx, sess.TaskId)
		if err != nil {
			return err
		}
		if !found || lead.Id != sess.Id {
			continue
		}
		task, err := m.store.GetTask(ctx, sess.TaskId)
		if err != nil {
			return err
		}
		merging, err := git.MergeInProgress(task.WorktreePath)
		if err != nil {
			return err
		}
		if merging {
			m.ownMerge(sess.Id, task.Id, task.WorktreePath)
		}
	}
	return nil
}

func (m *Manager) endMergeTurn(sessionID string, abort bool) {
	m.mu.Lock()
	t, found := m.mergeTurns[sessionID]
	delete(m.mergeTurns, sessionID)
	m.mu.Unlock()
	if !found {
		return
	}
	if abort {
		err := git.AbortMerge(t.dir)
		if err != nil {
			slog.Error("abort unfinished merge", "session", sessionID, "err", err)
		}
	}
	t.unlock()
}