byre/mac/Sources/Byre/Store.swift

import Foundation
import Observation
import SwiftData

let taskTabs: [TaskTab] = [.transcript, .board, .changes, .files, .terminal, .workflow, .project]

enum JumpTarget: Equatable {
    case summary(Int64)
    case question(String)
    case review(Int)
}

struct Jump: Equatable {
    let taskId: String
    let target: JumpTarget
    let n: Int
}

private let historyPage = 1000
private let leadPage = 50
private let routeHash = Re(#"^#/(task|project)/(\w+)(?:/(\w+))?$"#)

struct TaskState {
    enum Kind {
        case answer, merging, held, working, idle
    }

    let kind: Kind
    let label: String
    let title: String
}

private let workingTitles: [ByreTask.StatePayload: (label: String, title: String)] = [
    .running: ("Working", "The agent is mid-turn or its subagents are running"),
    .reviewPending: ("Reviewing", "A review of the turn's work is about to start"),
    .reviewing: ("Reviewing", "A review is running"),
    .discarding: ("Discarding", "The agent is writing the task summary, then byre removes the task"),
]

@MainActor
@Observable
final class Store {
    let api: API
    var me: User?
    var projects: [Project] = []
    var tasks: [String: [ByreTask]] = [:]
    var sessions: [String: [AgentSession]] = [:]
    var inbox: [InboxItem] = []
    private(set) var selectedTaskId: String?
    var pinned: Int?
    var connected = false
    private(set) var focus = Focus(tab: .transcript)
    var leadSessions: [AgentSession] = []
    let leadFeed: Feed
    var jump: Jump?
    @ObservationIgnored private var jumps = 0
    var composerFocus: String?
    var summaryReads: [String: Int64] = [:]
    private(set) var leadMore = false
    var todoBusy: [String: String] = [:]
    var pickedMembers: [String: String] = [:]
    private(set) var drafts: Drafts?
    let images: ImageDrafts
    let imageCache = ImageCache()
    let events = Listeners<ByreEvent>()
    let paneFocus: PaneFocus
    let remote: Remote

    private var taskFeeds: [String: Feed] = [:]
    @ObservationIgnored private var historyLoads: Set<ObjectIdentifier> = []
    @ObservationIgnored private var requestSeqs: [String: Int] = [:]
    @ObservationIgnored private var globalStream: EventStream?
    @ObservationIgnored private var leadStream: EventStream?
    @ObservationIgnored private var stopped = false
    @ObservationIgnored var changesMemory = ChangesMemory()

    init(api: API, images: ImageDrafts) {
        self.api = api
        self.images = images
        leadFeed = Feed(shown: { true })
        paneFocus = PaneFocus(api: api)
        remote = Remote(api: api)
    }

    func stop() async {
        stopped = true
        globalStream?.stop()
        leadStream?.stop()
        remote.stop()
        drafts?.flush()
        await paneFocus.leave()
    }

    func loadDrafts(from old: Store?) {
        drafts = old?.drafts ?? Drafts(origin: api.origin, login: me!.login)
        drafts!.prune(self)
    }

    var todoStarting: Set<String> {
        Set(todoBusy.keys)
    }

    func startTodo(_ projectId: String, _ title: String, discuss: Bool, _ overlays: Overlays) {
        let key = "\(projectId)\n\(title)"
        if todoBusy[key] != nil {
            return
        }
        todoBusy[key] = discuss ? "discuss" : "start"
        Task {
            do {
                _ = try await api.client.startTodoItem(path: .init(projectId: projectId), body: .json(.init(title: title, model: defaultModel, discuss: discuss))).created
            } catch {
                overlays.showAlert(errorText(error))
            }
            todoBusy[key] = nil
        }
    }

    func startSuggestion(_ projectId: String, _ title: String, brief: String, _ overlays: Overlays) {
        let key = "\(projectId)\n\(title)"
        if todoBusy[key] != nil {
            return
        }
        todoBusy[key] = "start"
        Task {
            do {
                _ = try await api.client.createTask(path: .init(projectId: projectId), body: .json(.init(title: title, prompt: brief, model: defaultModel))).created
            } catch {
                overlays.showAlert(errorText(error))
            }
            todoBusy[key] = nil
        }
    }

    func discussSuggestion(_ projectId: String, _ title: String, brief: String, _ overlays: Overlays) {
        let key = "\(projectId)\n\(title)"
        if todoBusy[key] != nil {
            return
        }
        todoBusy[key] = "discuss"
        let lead = currentLead()
        let live = lead != nil && lead!.state != .completed
        Task {
            do {
                _ = try await api.client.messageLead(body: .json(.init(text: "Let's discuss your suggestion \"\(title)\" before starting it:\n\n\(brief)", model: live ? nil : defaultModel))).accepted
            } catch {
                overlays.showAlert(errorText(error))
            }
            todoBusy[key] = nil
        }
    }

    var transcript: [ByreEvent] {
        selectedFeed?.events ?? []
    }

    var selectedFeed: Feed? {
        selectedTaskId.flatMap { taskFeeds[$0] }
    }

    func activeTasks() -> [(project: Project, task: ByreTask)] {
        projects.flatMap { project in
            (tasks[project.id] ?? []).filter { $0.status == .active }.map { (project, $0) }
        }
    }

    func findTask(_ id: String) -> ByreTask? {
        for list in tasks.values {
            if let task = list.first(where: { $0.id == id }) {
                return task
            }
        }
        return nil
    }

    func currentSession(_ taskId: String) -> AgentSession? {
        let list = (sessions[taskId] ?? []).filter { $0.member == (pickedMember(taskId) ?? "") }
        return list.last(where: { $0.state != .completed }) ?? list.last
    }

    func team(_ taskId: String) -> [AgentSession] {
        let list = sessions[taskId] ?? []
        var names: [String] = []
        for s in list where !s.member.isEmpty && !names.contains(s.member) {
            names.append(s.member)
        }
        return names.compactMap { name in
            let mine = list.filter { $0.member == name }
            return mine.last(where: { $0.state != .completed }) ?? mine.last
        }
    }

    func pickedMember(_ taskId: String) -> String? {
        let names = team(taskId).map(\.member)
        if let picked = pickedMembers[taskId], names.contains(picked) {
            return picked
        }
        return names.first
    }

    func taskBusy(_ taskId: String) -> Bool {
        (sessions[taskId] ?? []).contains { $0.state == .running }
    }

    func taskState(_ task: ByreTask) -> TaskState {
        if inbox.contains(where: { $0.taskId == task.id }) {
            return TaskState(kind: .answer, label: "Needs your answer", title: "The agent asked a question")
        }
        if task.merging {
            return TaskState(kind: .merging, label: "Merging", title: "The merge turn is running")
        }
        if task.mergingIntoParent {
            return TaskState(kind: .merging, label: "Merging", title: "Waiting for the parent task's agent to merge it")
        }
        if task.mergeQueuedAt != nil {
            return TaskState(kind: .held, label: "Merge held", title: "Waiting for the merge in progress in this project to finish")
        }
        if task.mergeAfterReview {
            return TaskState(kind: .held, label: "Merge after review", title: "Merges when the running review finishes with no findings")
        }
        if let working = workingTitles[task.state] {
            return TaskState(kind: .working, label: working.label, title: working.title)
        }
        return TaskState(kind: .idle, label: "Idle", title: "Waiting for a message")
    }

    func currentLead() -> AgentSession? {
        leadSessions.last(where: { $0.state != .completed }) ?? leadSessions.last
    }

    func snapshots() -> [Snapshot] {
        transcript.filter { $0.kind == "snapshot.created" }.map { try! decodePayload($0.payload, as: Snapshot.self) }
    }

    func pinSnapshot(_ n: Int?) {
        pinned = n
    }

    func jumpTo(_ taskId: String, _ target: JumpTarget) {
        jumps += 1
        jump = Jump(taskId: taskId, target: target, n: jumps)
        focusTask(taskId, .transcript)
    }

    func showQuestion(_ taskId: String) {
        guard let item = inbox.filter({ $0.taskId == taskId }).min(by: { $0.createdAt < $1.createdAt }) else {
            focusTask(taskId, .transcript)
            return
        }
        if let member = sessions[taskId]?.first(where: { $0.id == item.agentSessionId })?.member, !member.isEmpty {
            pickedMembers[taskId] = member
        }
        jumpTo(taskId, .question(item.id))
    }

    var defaultModel: String {
        me!.defaultModel
    }

    func loadMe() async throws {
        me = try await api.client.getMe().ok.body.json
    }

    func loadProjects() async throws {
        projects = try await api.client.listProjects().ok.body.json
        try await each(projects.map(\.id)) { id in
            try await self.loadTasks(id)
        }
    }

    private func claimRequest(_ key: String) -> () -> Bool {
        let seq = (requestSeqs[key] ?? 0) + 1
        requestSeqs[key] = seq
        return { [weak self] in
            self?.requestSeqs[key] == seq
        }
    }

    func loadTasks(_ projectId: String) async throws {
        let latest = claimRequest("tasks:\(projectId)")
        let list = try await api.client.listTasks(path: .init(projectId: projectId)).ok.body.json
        if !latest() {
            return
        }
        tasks[projectId] = list
        for t in list where t.status == .active {
            _ = taskFeed(t.id)
        }
        dropClosedFeeds()
        drafts?.prune(self)
        try await each(list.filter { $0.status == .active }.map(\.id)) { id in
            try await self.loadSessions(id)
        }
    }

    func loadSessions(_ taskId: String) async throws {
        let latest = claimRequest("sessions:\(taskId)")
        let list = try await api.client.listSessions(path: .init(taskId: taskId)).ok.body.json
        if !latest() {
            return
        }
        sessions[taskId] = list
    }

    func messageTask(_ taskId: String, _ text: String) async throws {
        let sessions = try await api.client.listSessions(path: .init(taskId: taskId)).ok.body.json
        if let live = sessions.last(where: { $0.state != .completed }) {
            _ = try await api.client.sendMessage(path: .init(sessionId: live.id), body: .json(.init(text: text))).accepted
        } else {
            _ = try await api.client.createSession(path: .init(taskId: taskId), body: .json(.init(prompt: text))).created
        }
    }

    func loadInbox() async throws {
        let latest = claimRequest("inbox")
        let items = try await api.client.listInbox().ok.body.json
        if !latest() {
            return
        }
        inbox = items
        drafts?.prune(self)
    }

    func loadLeadSessions() async throws {
        leadSessions = try await api.client.listLeadSessions().ok.body.json
    }

    func loadFocus() async throws {
        let latest = claimRequest("focus")
        let f = try await api.client.getFocus().ok.body.json
        if !latest() {
            return
        }
        applyFocus(f)
    }

    func applyFocus(_ f: Focus) {
        focus = f
        if f.taskId != selectedTaskId {
            selectTask(f.taskId)
        }
    }

    func focusedProject() -> Project? {
        projects.first { $0.id == focus.projectId }
    }

    func setFocus(taskId: String?? = .none, projectId: String?? = .none, tab: TaskTab? = nil) async throws {
        let task = taskId ?? focus.taskId
        let project = task != nil ? findTask(task!)?.projectId : (projectId ?? focus.projectId)
        var next = Focus(tab: tab ?? focus.tab)
        if let task, !task.isEmpty {
            next.taskId = task
        }
        if let project, !project.isEmpty {
            next.projectId = project
        }
        let latest = claimRequest("focus")
        applyFocus(next)
        let f = try await api.client.setFocus(body: .json(next)).ok.body.json
        if !latest() {
            return
        }
        applyFocus(f)
    }

    func focusTask(_ taskId: String, _ tab: TaskTab = .transcript) {
        spawn("set focus") {
            try await self.setFocus(taskId: taskId, tab: tab)
        }
    }

    func focusProject(_ projectId: String) {
        spawn("set focus") {
            try await self.setFocus(taskId: .some(nil), projectId: projectId, tab: .project)
        }
    }

    func focusTab(_ tab: TaskTab) {
        spawn("set focus") {
            try await self.setFocus(tab: tab)
        }
    }

    func route(_ hash: String) {
        let m = routeHash.exec(hash)
        let kind = m?[1] ?? nil
        let id = m?[2] ?? nil
        let name = m?[3] ?? nil
        let tab = taskTabs.first { $0.rawValue == name }
        if kind == "task", let id {
            if id != focus.taskId || (tab != nil && tab != focus.tab) {
                spawn("set focus") {
                    try await self.setFocus(taskId: id, tab: tab)
                }
            }
            return
        }
        let projectId = kind == "project" ? id : nil
        if focus.taskId != nil || projectId != focus.projectId {
            spawn("set focus") {
                try await self.setFocus(taskId: .some(nil), projectId: .some(projectId))
            }
        }
    }

    private func loadHistory(_ taskId: String, _ feed: Feed) async throws {
        var got: [ByreEvent] = []
        while true {
            let page = try await api.client.listTaskEvents(path: .init(taskId: taskId), query: .init(after: got.last?.seq ?? 0, limit: historyPage)).ok.body.json
            got += page
            if page.count < historyPage {
                break
            }
        }
        feed.load(got)
    }

    private func taskFeed(_ taskId: String) -> Feed {
        let feed = taskFeeds[taskId] ?? Feed(shown: { [weak self] in
            self?.selectedTaskId == taskId
        })
        taskFeeds[taskId] = feed
        let key = ObjectIdentifier(feed)
        if feed.ready || historyLoads.contains(key) {
            return feed
        }
        historyLoads.insert(key)
        Task {
            do {
                try await loadHistory(taskId, feed)
            } catch {
                log.error("load task history \(taskId, privacy: .public): \(errorText(error), privacy: .public)")
            }
            historyLoads.remove(key)
        }
        return feed
    }

    private func dropClosedFeeds() {
        for taskId in taskFeeds.keys {
            if taskId != selectedTaskId, let task = findTask(taskId), task.status != .active {
                taskFeeds[taskId] = nil
            }
        }
    }

    func transcriptLoading() -> Bool {
        selectedTaskId != nil && selectedFeed?.ready != true
    }

    func leadLoading() -> Bool {
        !leadFeed.ready
    }

    func lastVisits() -> [String: Double] {
        fetchRecords(FetchDescriptor<VisitRecord>()).reduce(into: [:]) { visits, v in
            visits[v.taskId] = max(visits[v.taskId] ?? 0, v.at)
        }
    }

    private func markVisited(_ taskId: String?) {
        guard let taskId else {
            return
        }
        let at = Date().timeIntervalSince1970 * 1000
        let records = fetchRecords(FetchDescriptor<VisitRecord>(predicate: #Predicate { $0.taskId == taskId }))
        if records.isEmpty {
            persistence.insert(VisitRecord(taskId: taskId, at: at))
        } else {
            records[0].at = at
            for record in records.dropFirst() {
                persistence.delete(record)
            }
        }
        saveRecords()
    }

    func loadSummaryReads() async throws {
        summaryReads = try await api.client.getSummaryReads().ok.body.json.additionalProperties
    }

    func unreadSummary(_ task: ByreTask) -> Bool {
        task.workflow == nil && task.summarySeq > (summaryReads[task.id] ?? 0)
    }

    var unreadSummaries: Int {
        tasks.values.joined().filter { $0.status == .active && unreadSummary($0) }.count
    }

    private func noteSummaryRead(_ taskId: String, _ seq: Int64) -> Bool {
        if (summaryReads[taskId] ?? 0) >= seq {
            return false
        }
        summaryReads[taskId] = seq
        return true
    }

    func markSummaryRead(_ taskId: String, _ seq: Int64) {
        if !noteSummaryRead(taskId, seq) {
            return
        }
        spawn("mark summary read") {
            _ = try await self.api.client.markSummaryRead(path: .init(taskId: taskId), body: .json(.init(seq: seq))).noContent
        }
    }

    func openSummary(_ task: ByreTask) {
        let seq = task.summarySeq
        if let ev = taskFeeds[task.id]?.events.first(where: { $0.seq == seq }) {
            pickMember(task.id, ev)
            jumpTo(task.id, .summary(seq))
            return
        }
        Task {
            do {
                let evs = try await api.client.listTaskEvents(path: .init(taskId: task.id), query: .init(after: seq - 1, limit: 1)).ok.body.json
                if let ev = evs.first(where: { $0.seq == seq }) {
                    pickMember(task.id, ev)
                }
            } catch {
                log.error("load summary event: \(errorText(error), privacy: .public)")
            }
            jumpTo(task.id, .summary(seq))
        }
    }

    private func pickMember(_ taskId: String, _ ev: ByreEvent) {
        if let member = str(ev.payload.value, "member"), !member.isEmpty {
            pickedMembers[taskId] = member
        }
    }

    func selectTask(_ taskId: String?) {
        markVisited(selectedTaskId)
        markVisited(taskId)
        selectedTaskId = taskId
        if let taskId {
            _ = taskFeed(taskId)
        }
        pinned = nil
        dropClosedFeeds()
    }

    private func fetchLeadTail() async -> EventPage? {
        while !stopped {
            do {
                return try await api.client.listLeadEvents(query: .init(turns: leadPage)).ok.body.json
            } catch {
                log.error("load lead history: \(errorText(error), privacy: .public)")
                try? await Task.sleep(for: .seconds(2))
            }
        }
        return nil
    }

    func startLeadStream() async {
        guard let page = await fetchLeadTail(), !stopped else {
            return
        }
        leadMore = page.more
        leadStream = EventStream(api: api, after: page.events.last?.seq ?? 0, lead: true, onEvent: { [weak self] ev in
            self?.leadFeed.push(ev)
        }, onState: { _ in }, onReady: { [weak self] in
            self?.leadFeed.load(page.events)
        })
    }

    func startGlobalStream(after: Int64, reconnected: @escaping () -> Void) {
        globalStream = EventStream(api: api, after: after, lead: false, onEvent: { [weak self] ev in
            self?.handleEvent(ev)
        }, onState: { [weak self] up in
            self?.connected = up
            if up {
                self?.paneFocus.reconnected()
                reconnected()
            }
        }, onReady: {})
    }

    func leadHasOlder() -> Bool {
        leadFeed.ready && leadMore
    }

    func loadOlderLead() async throws -> [ByreEvent] {
        guard let first = leadFeed.events.first else {
            return []
        }
        let page = try await api.client.listLeadEvents(query: .init(before: first.seq, turns: leadPage)).ok.body.json
        leadMore = page.more
        leadFeed.prepend(page.events)
        return page.events
    }

    func handleEvent(_ ev: ByreEvent) {
        if let taskId = ev.taskId {
            taskFeeds[taskId]?.push(ev)
        }
        if ev.seq == 0 {
            events.send(ev)
            return
        }

        switch ev.kind {
        case "user.updated":
            if ev.payload.value["login"] as? String == me?.login {
                spawn("load me") {
                    try await self.loadMe()
                }
            }
        case "project.created":
            spawn("load projects") {
                try await self.loadProjects()
            }
        case "task.created", "task.merge_title", "task.merge_queued", "task.merge_unqueued", "task.merge_after_review", "task.merge_after_review_ended", "task.merge_after_review_cancelled", "task.parent_merge_started", "task.parent_merge_stopped", "task.merge_started", "task.merge_ended", "task.discard_started", "task.merged", "task.discarded", "task.state":
            if let projectId = ev.projectId {
                spawn("load tasks") {
                    try await self.loadTasks(projectId)
                }
            }
        case "session.created", "session.state", "session.subagents", "session.model":
            if ev.owner != nil {
                spawn("load lead sessions") {
                    try await self.loadLeadSessions()
                }
            } else if let taskId = ev.taskId {
                spawn("load sessions") {
                    try await self.loadSessions(taskId)
                }
            }
        case "inbox.created":
            notifyInbox(self, try! decodePayload(ev.payload, as: InboxItem.self))
            spawn("load inbox") {
                try await self.loadInbox()
            }
        case "inbox.updated":
            clearInbox(try! decodePayload(ev.payload, as: InboxItem.self))
            spawn("load inbox") {
                try await self.loadInbox()
            }
        case "session.error", "session.rate_limited":
            notifyError(self, ev)
        case "turn.ended":
            notifyTurn(self, ev)
            if ev.payload.value["summary"] as? Bool == true, let taskId = ev.taskId {
                updateTask(taskId) { task in
                    if ev.seq > task.summarySeq {
                        task.summarySeq = ev.seq
                    }
                }
            }
        case "summary.read":
            if let taskId = ev.payload.value["task_id"] as? String, let seq = ev.payload.value["seq"] as? Int {
                _ = noteSummaryRead(taskId, Int64(seq))
                clearTurns(taskId, through: Int64(seq))
            }
        case "focus.updated":
            if ev.actor != me?.login {
                _ = claimRequest("focus")
                applyFocus(try! decodePayload(ev.payload, as: Focus.self))
            }
        default:
            break
        }

        events.send(ev)
    }

    private func updateTask(_ id: String, _ change: (inout ByreTask) -> Void) {
        for (projectId, list) in tasks {
            if let i = list.firstIndex(where: { $0.id == id }) {
                change(&tasks[projectId]![i])
            }
        }
    }

    func spawn(_ what: String, _ op: @escaping @MainActor () async throws -> Void) {
        Task {
            do {
                try await op()
            } catch {
                log.error("\(what, privacy: .public): \(errorText(error), privacy: .public)")
            }
        }
    }

    private func each(_ ids: [String], _ op: @escaping @Sendable @MainActor (String) async throws -> Void) async throws {
        try await withThrowingTaskGroup(of: Void.self) { group in
            for id in ids {
                group.addTask {
                    try await op(id)
                }
            }
            try await group.waitForAll()
        }
    }
}