byre/mac/Sources/Byre/WorkflowEditor.swift

import ByreAPI
import SwiftUI

typealias WorkflowDefinition = Components.Schemas.WorkflowDefinition
typealias WorkflowStep = Components.Schemas.WorkflowStep
typealias TeamMember = Components.Schemas.TeamMember

@MainActor
@Observable
final class StepNode: Identifiable {
    enum Kind {
        case stage, order, parallel
    }

    let id = UUID()
    var kind: Kind
    var name = ""
    var model = ""
    var prompt = ""
    var agents: [TeamMember]?
    var children: [StepNode] = []
    weak var parent: StepNode?
    var open = false

    init(kind: Kind, children: [StepNode] = []) {
        self.kind = kind
        self.children = children
        for c in children {
            c.parent = self
        }
    }

    convenience init(_ step: WorkflowStep) {
        if let order = step.order {
            self.init(kind: .order, children: order.map { StepNode($0) })
        } else if let parallel = step.parallel {
            self.init(kind: .parallel, children: parallel.map { StepNode($0) })
        } else {
            self.init(kind: .stage)
            name = step.stage ?? ""
            model = step.model ?? ""
            prompt = step.prompt ?? ""
            agents = step.agents
        }
    }

    var step: WorkflowStep {
        switch kind {
        case .stage:
            WorkflowStep(stage: name, model: agents == nil ? model : nil, prompt: prompt, agents: agents)
        case .order:
            WorkflowStep(order: children.map(\.step))
        case .parallel:
            WorkflowStep(parallel: children.map(\.step))
        }
    }

    var index: Int {
        parent!.children.firstIndex { $0 === self }!
    }

    func contains(_ other: StepNode) -> Bool {
        other === self || children.contains { $0.contains(other) }
    }

    func stageNames() -> [String] {
        kind == .stage ? [name] : children.flatMap { $0.stageNames() }
    }

    func insert(_ node: StepNode, at i: Int) {
        node.parent = self
        children.insert(node, at: i)
    }

    func detach() {
        parent?.children.remove(at: index)
        parent = nil
    }
}

@MainActor
private func newStage(_ root: StepNode, _ model: String) -> StepNode {
    let names = Set(root.stageNames())
    var n = names.count + 1
    while names.contains("Stage \(n)") {
        n += 1
    }
    let s = StepNode(kind: .stage)
    s.name = "Stage \(n)"
    s.model = model
    s.open = true
    return s
}

@MainActor
private func stepGraph(_ root: StepNode) -> Flowchart {
    var chart = Flowchart(direction: .td)
    func walk(_ node: StepNode, _ before: [String]) -> [String] {
        switch node.kind {
        case .order:
            return node.children.reduce(before) { walk($1, $0) }
        case .parallel:
            return node.children.flatMap { walk($0, before) }
        case .stage:
            let id = "n\(chart.order.count)"
            chart.node(id, label: node.name.isEmpty ? " " : node.name, shape: .box)
            for b in before {
                chart.edges.append(FlowEdge(from: b, to: id, label: nil, style: .solid, arrow: true))
            }
            return [id]
        }
    }
    _ = walk(root, [])
    return chart
}

private struct StepLine: Identifiable {
    let node: StepNode
    let guide: String
    var id: UUID { node.id }
}

@MainActor
private func stepLines(_ node: StepNode, _ lead: String) -> [StepLine] {
    var out: [StepLine] = []
    for (i, c) in node.children.enumerated() {
        let last = i == node.children.count - 1
        out.append(StepLine(node: c, guide: lead + (last ? "└─ " : "├─ ")))
        out += stepLines(c, lead + (last ? "   " : "│  "))
    }
    return out
}

struct WorkflowEditor: View {
    let original: String?
    let close: () -> Void
    @Environment(Store.self) private var store
    @Environment(Overlays.self) private var overlays
    @ViewState private var root: StepNode?
    @ViewState private var loaded: WorkflowDefinition?
    @ViewState private var name = ""
    @ViewState private var description = ""
    @ViewState private var failure = ""
    @ViewState private var saving = false

    var body: some View {
        ScrollView {
            VStack(alignment: .leading, spacing: 0) {
                HStack(spacing: 10) {
                    Button("‹ workflows") {
                        Task {
                            await leave()
                        }
                    }
                    .buttonStyle(SmallLinkButtonStyle())
                    Spacer()
                    Button("save") {
                        Task {
                            await save()
                        }
                    }
                    .buttonStyle(WebButtonStyle(primary: true))
                    .disabled(saving || !changed || target.isEmpty)
                }
                .padding(.top, 16)
                if let root {
                    TextField("Name", text: $name)
                        .modifier(FieldStyle())
                        .padding(.top, 12)
                    TextField("Description", text: $description, axis: .vertical)
                        .modifier(FieldStyle())
                        .padding(.top, 8)
                    if !root.stageNames().isEmpty {
                        FlowchartCanvas(layout: layoutFlowchart(stepGraph(root)))
                            .frame(maxWidth: .infinity, alignment: .leading)
                            .padding(.top, 12)
                    }
                    if !failure.isEmpty {
                        Text(failure)
                            .foregroundStyle(Palette.danger)
                            .textSelection(.enabled)
                            .fixedSize(horizontal: false, vertical: true)
                            .padding(.top, 8)
                    }
                    VStack(alignment: .leading, spacing: 0) {
                        ForEach(stepLines(root, "")) { row in
                            StepRow(node: row.node, guide: row.guide, root: root)
                        }
                    }
                    .padding(.top, 12)
                    Button("add stage") {
                        root.insert(newStage(root, store.defaultModel), at: root.children.count)
                    }
                    .buttonStyle(SmallLinkButtonStyle())
                    .padding(.top, 8)
                } else if !failure.isEmpty {
                    Text(failure).foregroundStyle(Palette.muted).padding(.top, 12)
                }
            }
            .padding(.horizontal, 16)
            .padding(.bottom, 16)
            .frame(maxWidth: .infinity, alignment: .leading)
        }
        .frame(maxWidth: .infinity, maxHeight: .infinity)
        .background(Palette.bg)
        .task {
            await load()
        }
    }

    private var current: WorkflowDefinition? {
        root.map { WorkflowDefinition(description: description.isEmpty ? nil : description, steps: $0.children.map(\.step)) }
    }

    private var target: String {
        name.trimmingCharacters(in: .whitespaces)
    }

    private var changed: Bool {
        original == nil || current != loaded || target != original
    }

    private func load() async {
        guard let original else {
            let r = StepNode(kind: .order)
            r.insert(newStage(r, store.defaultModel), at: 0)
            root = r
            return
        }
        do {
            let def = try await store.api.client.getWorkflow(path: .init(name: original)).ok.body.json
            name = original
            description = def.description ?? ""
            root = StepNode(kind: .order, children: def.steps.map { StepNode($0) })
            loaded = current
        } catch {
            failure = errorText(error)
        }
    }

    private func save() async {
        guard let def = current else {
            return
        }
        saving = true
        defer { saving = false }
        do {
            _ = try await store.api.client.putWorkflow(path: .init(name: target), query: .init(from: original), body: .json(def)).ok.body.json
            close()
        } catch {
            failure = errorText(error)
        }
    }

    private func leave() async {
        if changed, current != nil, !(await overlays.showConfirm(["This workflow has unsaved changes."], action: "Discard changes", danger: true)) {
            return
        }
        close()
    }
}

private struct StepRow: View {
    @Bindable var node: StepNode
    let guide: String
    let root: StepNode
    @Environment(Store.self) private var store
    @ViewState private var target = false

    var body: some View {
        VStack(alignment: .leading, spacing: 6) {
            HStack(alignment: .firstTextBaseline, spacing: 6) {
                Text(guide).foregroundStyle(Palette.muted)
                switch node.kind {
                case .stage:
                    Button {
                        node.open.toggle()
                    } label: {
                        Image(systemName: node.open ? "chevron.down" : "chevron.right")
                            .font(.system(size: 9, weight: .semibold))
                            .foregroundStyle(Palette.muted)
                            .frame(width: 12)
                    }
                    .buttonStyle(.plain)
                    TextField("Stage name", text: $node.name)
                        .textFieldStyle(.plain)
                        .font(Fonts.font(13, weight: 700))
                    if let agents = node.agents {
                        Text("team of \(agents.count)").foregroundStyle(Palette.muted).fixedSize()
                    } else {
                        StepModelMenu(model: $node.model)
                    }
                case .order:
                    Text("in order").foregroundStyle(Palette.muted)
                    Spacer()
                case .parallel:
                    Text("in parallel").foregroundStyle(Palette.muted)
                    Spacer()
                }
                actions
            }
            .padding(.vertical, 3)
            .background(target ? Palette.accentTint : .clear, in: RoundedRectangle(cornerRadius: 4))
            .draggable(node.id.uuidString)
            .dropDestination(for: String.self) { items, _ in
                drop(items)
            } isTargeted: {
                target = $0
            }
            if node.kind == .stage && node.open {
                VStack(alignment: .leading, spacing: 6) {
                    if node.agents != nil {
                        Text("Every agent's prompt").foregroundStyle(Palette.muted)
                    }
                    PromptEditor(text: $node.prompt)
                    if let agents = node.agents {
                        ForEach(agents.indices, id: \.self) { i in
                            HStack(spacing: 6) {
                                TextField("Agent name", text: Binding(get: { node.agents?[i].name ?? "" }, set: { node.agents?[i].name = $0 }))
                                    .textFieldStyle(.plain)
                                    .font(Fonts.font(13, weight: 700))
                                StepModelMenu(model: Binding(get: { node.agents?[i].model ?? "" }, set: { node.agents?[i].model = $0 }))
                            }
                            .padding(.top, 4)
                            PromptEditor(text: Binding(get: { node.agents?[i].prompt ?? "" }, set: { node.agents?[i].prompt = $0 }))
                        }
                    }
                }
                .padding(.leading, CGFloat(guide.count) * 8 + 18)
                .padding(.bottom, 8)
            }
        }
    }

    private var actions: some View {
        Menu {
            switch node.kind {
            case .stage:
                Button("Add stage after") { addAfter() }
                Button("Run in parallel with a new stage") { wrap(.parallel) }
                Button("Group with a new stage in order") { wrap(.order) }
            case .order, .parallel:
                Button("Add stage inside") { node.insert(stage(), at: node.children.count) }
                Button("Add stage after") { addAfter() }
                Button(node.kind == .order ? "Change to in parallel" : "Change to in order") {
                    node.kind = node.kind == .order ? .parallel : .order
                }
                Button("Ungroup") { ungroup() }
            }
            Divider()
            let siblings = node.parent?.children ?? []
            let i = siblings.firstIndex { $0 === node }
            Button("Move up") { move(-1) }
                .disabled(i == nil || i == 0)
            Button("Move down") { move(1) }
                .disabled(i == nil || i == siblings.count - 1)
            Divider()
            Button("Remove") { node.detach() }
        } label: {
            Image(systemName: "ellipsis")
                .font(.system(size: 11, weight: .semibold))
                .foregroundStyle(Palette.muted)
                .frame(width: 20, height: 16)
                .contentShape(Rectangle())
        }
        .menuStyle(.button)
        .buttonStyle(.plain)
        .menuIndicator(.hidden)
        .fixedSize()
    }

    private func stage() -> StepNode {
        newStage(root, store.defaultModel)
    }

    private func addAfter() {
        node.parent!.insert(stage(), at: node.index + 1)
    }

    private func wrap(_ kind: StepNode.Kind) {
        let parent = node.parent!
        let i = node.index
        node.detach()
        let group = StepNode(kind: kind)
        group.insert(node, at: 0)
        group.insert(stage(), at: 1)
        parent.insert(group, at: i)
    }

    private func ungroup() {
        let parent = node.parent!
        let i = node.index
        let children = node.children
        node.detach()
        for (n, c) in children.enumerated() {
            parent.insert(c, at: i + n)
        }
    }

    private func move(_ by: Int) {
        let parent = node.parent!
        let i = node.index
        node.detach()
        parent.insert(node, at: i + by)
    }

    private func drop(_ items: [String]) -> Bool {
        guard let id = items.first.flatMap(UUID.init(uuidString:)), let dragged = find(root, id), !dragged.contains(node) else {
            return false
        }
        dragged.detach()
        if node.kind == .stage {
            node.parent!.insert(dragged, at: node.index)
        } else {
            node.insert(dragged, at: node.children.count)
        }
        return true
    }

    private func find(_ n: StepNode, _ id: UUID) -> StepNode? {
        if n.id == id {
            return n
        }
        for c in n.children {
            if let f = find(c, id) {
                return f
            }
        }
        return nil
    }
}

private struct StepModelMenu: View {
    @Binding var model: String

    var body: some View {
        Menu {
            ForEach(models, id: \.self) { m in
                Button(m) { model = m }
            }
        } label: {
            HStack(spacing: 6) {
                Text(model.isEmpty ? "model" : model)
                Image(systemName: "chevron.up.chevron.down")
                    .font(.system(size: 8, weight: .semibold))
            }
            .font(Fonts.font(12))
            .foregroundStyle(model.isEmpty ? Palette.muted : Palette.text)
            .contentShape(Rectangle())
        }
        .menuStyle(.button)
        .buttonStyle(.plain)
        .menuIndicator(.hidden)
        .fixedSize()
    }
}

private struct PromptEditor: View {
    @Binding var text: String
    @FocusState private var focused: Bool

    var body: some View {
        TextEditor(text: $text)
            .font(Fonts.font(12))
            .scrollContentBackground(.hidden)
            .writingToolsBehavior(.disabled)
            .focused($focused)
            .padding(.horizontal, 5)
            .padding(.vertical, 4)
            .frame(height: 160)
            .background(Palette.raised, in: RoundedRectangle(cornerRadius: 6))
            .overlay(alignment: .topLeading) {
                if text.isEmpty {
                    Text("Prompt…")
                        .foregroundStyle(Palette.muted)
                        .padding(.horizontal, 10)
                        .padding(.vertical, 4)
                        .allowsHitTesting(false)
                }
            }
            .overlay(RoundedRectangle(cornerRadius: 6).stroke(focused ? Palette.accent : Palette.border))
    }
}