import ByreAPI
import SwiftUI

typealias WorkflowRun = Components.Schemas.WorkflowRun
typealias RunStage = Components.Schemas.WorkflowRunStage
typealias StageAttempt = Components.Schemas.StageAttempt

private let stateTints: [RunStage.StatePayload: BoxTint] = [
    .waiting: .plain,
    .running: .blue,
    .idle: .plain,
    .finished: .green,
    .merged: .green,
    .cleanedUp: .green,
    .discarded: .red,
]

private func stageSteps(_ steps: [WorkflowStep]) -> [WorkflowStep] {
    steps.flatMap { s in s.order.map(stageSteps) ?? s.parallel.map(stageSteps) ?? [s] }
}

private func runDiagram(_ run: WorkflowRun) -> String {
    let steps = Dictionary(stageSteps(run.definition.steps).map { ($0.stage ?? "", $0) }, uniquingKeysWith: { a, _ in a })
    let index = Dictionary(run.stages.enumerated().map { ($0.element.name, $0.offset) }, uniquingKeysWith: { a, _ in a })
    return stageDiagram(run.stages.map { s in
        GraphStage(name: s.name, tint: stateTints[s.state]!, agents: steps[s.name]?.agents, after: s.after.compactMap { index[$0] })
    })
}

private func attemptLabel(_ a: StageAttempt) -> String {
    switch a.state {
    case .running: "running"
    case .idle: "idle, not finished"
    case .finished: "finished"
    case .merged: "merged"
    case .discarded: "discarded before finishing\(a.discardedBy.map { " by \($0)" } ?? "")"
    case .cleanedUp: a.discardedBy == "byre" ? "finished, then discarded by the run's cleanup" : "finished, then discarded\(a.discardedBy.map { " by \($0)" } ?? "")"
    }
}

private func stageLabel(_ s: RunStage) -> String {
    if s.state != .waiting {
        return attemptLabel(s.attempts.last { $0.output != nil } ?? s.attempts.last!)
    }
    if !s.waitingOn.isEmpty {
        return "waiting on \(s.waitingOn.map { "“\($0)”" }.joined(separator: ", "))"
    }
    return "waiting to start"
}

@MainActor
private func taskLink(_ taskId: String, _ label: String, bold: Bool = false) -> AttributedString {
    var a = AttributedString(label)
    a.link = AppLink.task(taskId)
    a.foregroundColor = Palette.accent
    a.font = Fonts.font(13, weight: bold ? 700 : 400)
    return a
}

private func stateColor(_ state: String) -> Color {
    switch state {
    case "running": Palette.accent
    case "idle", "discarding": Palette.muted
    case "finished", "merged": Palette.ok
    default: Palette.muted
    }
}

struct WorkflowView: View {
    @Environment(Store.self) private var store
    @ViewState private var run: WorkflowRun?
    @ViewState private var runId: String?
    @ViewState private var error = ""
    @ViewState private var target = 0
    @ViewState private var loading = false
    @ViewState private var again = false

    var body: some View {
        let task = store.selectedTaskId.flatMap { store.findTask($0) }
        VStack(spacing: 0) {
            if let task, task.workflow != nil {
                if let run, runId == task.id {
                    HStack(alignment: .firstTextBaseline, spacing: 10) {
                        Text(run.name).font(Fonts.font(13, weight: 700))
                        Text(run.title).foregroundStyle(Palette.muted)
                        Text(run.finished ? "workflow finished" : "in progress")
                            .font(Fonts.font(11))
                            .foregroundStyle(run.finished ? Palette.ok : Palette.accent)
                    }
                    .modifier(HeaderBar())
                    ScrollView {
                        VStack(alignment: .leading, spacing: 12) {
                            BoxDiagramView(source: runDiagram(run))
                            if !error.isEmpty {
                                EmptyText(text: error, error: true)
                            } else {
                                if !trimmed(run.brief).isEmpty {
                                    DetailsBox(boxed: false) {
                                        Text("Brief").foregroundStyle(Palette.muted)
                                    } content: {
                                        MarkdownView(blocks: parseMarkdown(run.brief), ctx: MDContext(taskId: task.id))
                                            .padding(.top, 6)
                                    }
                                }
                                ForEach(run.stages, id: \.name) { s in
                                    StageCard(run: run, stage: s)
                                }
                            }
                        }
                        .padding(.horizontal, 16)
                        .padding(.top, 8)
                        .padding(.bottom, 16)
                    }
                } else if !error.isEmpty {
                    EmptyText(text: error, error: true)
                } else {
                    EmptyText(text: "Loading…")
                }
            } else {
                EmptyText(text: task != nil ? "This task isn't a workflow run." : "No task selected.")
            }
        }
        .frame(maxHeight: .infinity, alignment: .top)
        .task(id: store.selectedTaskId) {
            target += 1
            run = nil
            runId = nil
            error = ""
            await load()
            let off = store.events.add { ev in
                if relevant(ev) {
                    Task {
                        await load()
                    }
                }
            }
            await untilCancelled()
            off()
        }
        .onChange(of: task?.workflow) {
            if run == nil && !loading && error.isEmpty {
                Task {
                    await load()
                }
            }
        }
    }

    private func relevant(_ ev: ByreEvent) -> Bool {
        guard ev.seq != 0, let run, let runId else {
            return false
        }
        if ev.kind == "task.created" {
            return ev.payload.value["parent_task_id"] as? String == runId
        }
        let ids = Set(run.stages.flatMap { $0.attempts.map(\.taskId) } + [runId])
        guard let taskId = ev.taskId else {
            return false
        }
        return ids.contains(taskId) && (ev.kind.hasPrefix("task.") || ev.kind.hasPrefix("session."))
    }

    private func load() async {
        guard let taskId = store.selectedTaskId, store.findTask(taskId)?.workflow != nil else {
            return
        }
        if loading {
            again = true
            return
        }
        loading = true
        target += 1
        let started = target
        do {
            let fresh = try await store.api.client.getWorkflowRun(path: .init(taskId: taskId)).ok.body.json
            if started == target {
                run = fresh
                runId = taskId
                error = ""
            }
        } catch {
            if started == target {
                self.error = errorText(error)
            }
        }
        loading = false
        if again {
            again = false
            await load()
        }
    }
}

private struct StageCard: View {
    let run: WorkflowRun
    let stage: RunStage

    var body: some View {
        let only = stage.attempts.count == 1 ? stage.attempts[0] : nil
        VStack(alignment: .leading, spacing: 2) {
            HStack(alignment: .firstTextBaseline, spacing: 10) {
                if let only {
                    Text(taskLink(only.taskId, stage.name, bold: true))
                } else {
                    Text(stage.name).font(Fonts.font(13, weight: 700))
                }
                Text(stageLabel(stage))
                    .font(Fonts.font(11))
                    .foregroundStyle(stateColor(stage.state.rawValue))
                if only?.discarding == true {
                    Text("discarding…").font(Fonts.font(11)).foregroundStyle(Palette.muted)
                }
            }
            if let only, only.title.caseInsensitiveCompare(stage.name) != .orderedSame {
                Text("task “\(only.title)”")
                    .font(Fonts.font(11))
                    .foregroundStyle(Palette.muted)
            }
            if !stage.after.isEmpty {
                Text("after \(stage.after.map { "“\($0)”" }.joined(separator: ", "))")
                    .font(Fonts.font(11))
                    .foregroundStyle(Palette.muted)
            }
            if let only {
                AttemptRow(run: run, stage: stage, attempt: only, index: 0, solo: true)
            } else if !stage.attempts.isEmpty {
                VStack(alignment: .leading, spacing: 8) {
                    ForEach(Array(stage.attempts.enumerated()), id: \.offset) { i, a in
                        AttemptRow(run: run, stage: stage, attempt: a, index: i, solo: false)
                    }
                }
                .padding(.top, 8)
            }
        }
        .padding(.horizontal, 14)
        .padding(.vertical, 10)
        .frame(maxWidth: .infinity, alignment: .leading)
        .background(Palette.raised, in: RoundedRectangle(cornerRadius: 10))
        .overlay(RoundedRectangle(cornerRadius: 10).stroke(Palette.border))
    }
}

private struct AttemptRow: View {
    let run: WorkflowRun
    let stage: RunStage
    let attempt: StageAttempt
    let index: Int
    let solo: Bool
    @Environment(Store.self) private var store

    var body: some View {
        let a = attempt
        let fed = run.stages.flatMap { s in s.attempts.filter { d in d.inputs.contains { $0.taskId == a.taskId } }.map(\.taskId) }
        VStack(alignment: .leading, spacing: 2) {
            if !solo {
                HStack(alignment: .firstTextBaseline, spacing: 10) {
                    Text("#\(index + 1)").foregroundStyle(Palette.muted)
                    Text(taskLink(a.taskId, a.title, bold: a.taskId == store.selectedTaskId))
                    Text(attemptLabel(a))
                        .font(Fonts.font(11))
                        .foregroundStyle(stateColor(a.state.rawValue))
                    if a.discarding {
                        Text("discarding…").font(Fonts.font(11)).foregroundStyle(Palette.muted)
                    }
                }
            }
            if let from = a.forkedFrom {
                line([("fork of ", nil), (attemptName(from), from)])
            }
            if !a.inputs.isEmpty {
                line([("started from ", nil)] + joined(a.inputs.map(\.taskId)))
            }
            if !fed.isEmpty {
                line([("fed ", nil)] + joined(fed))
            }
            if let output = a.output {
                DetailsBox(boxed: false) {
                    HStack(spacing: 0) {
                        Text("Stage output").foregroundStyle(Palette.muted)
                        if let commit = a.commit {
                            Text(" \(String(commit.prefix(10)))").foregroundStyle(Palette.muted)
                        }
                    }
                } content: {
                    MarkdownView(blocks: parseMarkdown(output), ctx: MDContext(taskId: a.taskId))
                        .padding(.leading, 10)
                        .overlay(alignment: .leading) {
                            Rectangle().fill(Palette.border).frame(width: 1)
                        }
                        .padding(.top, 6)
                }
                .padding(.top, 6)
            }
        }
        .padding(.top, solo ? 0 : 8)
        .overlay(alignment: .top) {
            if !solo {
                Rectangle().fill(Palette.border).frame(height: 1)
            }
        }
    }

    private func joined(_ ids: [String]) -> [(String, String?)] {
        var out: [(String, String?)] = []
        for (n, id) in ids.enumerated() {
            if n > 0 {
                out.append((", ", nil))
            }
            out.append((attemptName(id), id))
        }
        return out
    }

    private func attemptName(_ taskId: String) -> String {
        for s in run.stages {
            if let i = s.attempts.firstIndex(where: { $0.taskId == taskId }) {
                return s.attempts.count > 1 ? "“\(s.name)” attempt \(i + 1)" : "“\(s.name)”"
            }
        }
        return store.findTask(taskId)?.title ?? taskId
    }

    private func line(_ parts: [(String, String?)]) -> some View {
        var out = AttributedString()
        for (text, taskId) in parts {
            if let taskId {
                var a = taskLink(taskId, text)
                a.font = Fonts.font(11)
                out += a
            } else {
                out += AttributedString(text)
            }
        }
        return Text(out)
            .font(Fonts.font(11))
            .foregroundStyle(Palette.muted)
            .padding(.top, 2)
    }
}