import Foundation
import Observation

enum StemServerStatus: Equatable {
    case unknown
    case online(tracks: Int, transcode: Bool)
    case offline(reason: String, cachedTracks: Int)

    var isOnline: Bool { if case .online = self { true } else { false } }
}

private struct ServerInfo: Decodable {
    var tracks = 0
    var ffmpeg = false
    init(from decoder: Decoder) throws {
        let c = try decoder.container(keyedBy: AnyKey.self)
        tracks = c.int("tracks") ?? 0
        ffmpeg = c.bool("ffmpeg") ?? false
    }
}

/// What stemd is doing, from stemserve's /api/queue.
enum StemState: Equatable { case stemming, queued, failed }

private struct QueueInfo: Decodable {
    var pending: [String] = []
    var current: String?
    var done = 0
    var failed: [String] = []
    init(from decoder: Decoder) throws {
        let c = try decoder.container(keyedBy: AnyKey.self)
        pending = c.list("pending")
        current = c.str("current")
        done = c.int("done") ?? 0
        failed = c.list("failed")
    }
}

/// Client for server/stemserve.py: the stems index, stem file URLs and waveforms.
@Observable
final class StemServer: @unchecked Sendable {
    @ObservationIgnored private let http: Http
    @ObservationIgnored private let settings: Settings
    @ObservationIgnored private let indexFile: URL
    @ObservationIgnored private let etagFile: URL
    @ObservationIgnored private let refreshLock = AsyncLock()
    @ObservationIgnored private var transcode = false
    @ObservationIgnored private var indexListeners: [() -> Void] = []
    private static let checkSeconds: UInt64 = 60

    private(set) var index = StemIndex([])
    private(set) var status = StemServerStatus.unknown
    /// When the index was last checked against the server, nil before the first check.
    private(set) var checkedAt: Date?
    /// stemd's live state: the song being separated, the ones waiting, the ones that failed.
    private(set) var current: String?
    private(set) var pending: Set<String> = []
    private(set) var failed: Set<String> = []
    @ObservationIgnored private var lastDone: Int?

    var busy: Bool { current != nil || !pending.isEmpty }

    init(http: Http, settings: Settings, dir: URL) {
        self.http = http
        self.settings = settings
        indexFile = dir.appendingPathComponent("stem_index.json")
        etagFile = dir.appendingPathComponent("stem_index.etag")
        if let data = try? Data(contentsOf: indexFile),
           let parsed = try? JSONDecoder().decode(StemIndexFile.self, from: data) {
            index = StemIndex(parsed.tracks, roots: roots())
        }
        // stemd on the server separates new music around the clock; look for its new stems
        // every minute. Unchanged indexes cost one 304.
        Task.detached { [weak self] in
            while let self {
                await self.refresh()
                try? await Task.sleep(nanoseconds: Self.checkSeconds * 1_000_000_000)
            }
        }
        // Watch stemd's queue: every few seconds while it has work, slowly when idle.
        Task.detached { [weak self] in
            while let self {
                await self.pollQueue()
                try? await Task.sleep(nanoseconds: (self.busy ? 4 : 30) * 1_000_000_000)
            }
        }
    }

    /// Where a track stands with stemd, nil when it has stems or isn't known to it.
    func state(_ track: Track) -> StemState? {
        guard let f = track.partFile, !f.isEmpty else { return nil }
        if find(track) != nil { return nil }
        if let c = current, f == c || f.hasSuffix("/" + c) { return .stemming }
        if pending.contains(f) { return .queued }
        if failed.contains(f) { return .failed }
        return nil
    }

    /// "Stemming “X” · N queued", nil when stemd is idle.
    var liveLine: String? {
        let n = pending.count
        if let c = current {
            let name = ((c as NSString).lastPathComponent as NSString).deletingPathExtension
            return "Stemming “\(name)”" + (n > 0 ? " · \(n) queued" : "")
        }
        return n > 0 ? "\(n) queued for stems" : nil
    }

    /// Reads /api/queue. A change in the finished count pulls the new index straight away.
    func pollQueue() async {
        guard !base.isEmpty,
              let u = try? url("/api/queue"),
              let data = try? await http.data(req(u)),
              let q = try? JSONDecoder().decode(QueueInfo.self, from: data) else { return }
        let changed = lastDone != nil && lastDone != q.done
        lastDone = q.done
        await MainActor.run {
            self.current = q.current
            self.pending = Set(q.pending)
            self.failed = Set(q.failed)
        }
        if changed { await refresh() }
    }

    private var base: String { settings.stemUrl.trimmingCharacters(in: CharacterSet(charactersIn: "/")) }
    private func roots() -> [String] { settings.stemRoots.split(separator: "\n").map(String.init).filter { !$0.isEmpty } }

    /// Main-thread callback after the index is replaced.
    func onIndexChange(_ f: @escaping () -> Void) { indexListeners.append(f) }

    private func req(_ url: URL, method: String = "GET") -> URLRequest {
        var r = URLRequest(url: url)
        r.httpMethod = method
        let t = settings.stemToken
        if !t.isEmpty { r.setValue(t, forHTTPHeaderField: "X-Stem-Token") }
        return r
    }

    private func url(_ path: String, _ params: [(String, Any?)] = []) throws -> URL {
        guard let u = URL(string: base + path) else { throw PlainError("Bad stemserve address") }
        return u.with(params)
    }

    func refresh(force: Bool = false) async {
        await refreshLock.run {
            do {
                let info = try JSONDecoder().decode(ServerInfo.self, from: try await http.data(req(try url("/api/info"))))
                transcode = info.ffmpeg
                let fm = FileManager.default
                let etag = (!force && fm.fileExists(atPath: indexFile.path)) ? try? String(contentsOf: etagFile, encoding: .utf8) : nil
                var r = req(try url("/api/index"))
                if let etag { r.setValue(etag, forHTTPHeaderField: "If-None-Match") }
                // Straight to the session: a 304 must not be confused with a cached 200.
                let (body, resp) = try await http.session.data(for: r)
                let code = (resp as? HTTPURLResponse)?.statusCode ?? 0
                var newIndex: StemIndex?
                if code == 304 {
                } else if (200..<300).contains(code) {
                    let parsed = try JSONDecoder().decode(StemIndexFile.self, from: body)
                    newIndex = StemIndex(parsed.tracks, roots: roots())
                    try? body.write(to: indexFile, options: .atomic)
                    if let tag = (resp as? HTTPURLResponse)?.value(forHTTPHeaderField: "ETag") {
                        try? tag.write(to: etagFile, atomically: true, encoding: .utf8)
                    }
                } else {
                    throw HttpError(code: code, message: "stemserve answered \(code)")
                }
                let fresh = newIndex
                await MainActor.run {
                    if let newIndex = fresh {
                        self.index = newIndex
                        self.indexListeners.forEach { $0() }
                    }
                    self.status = .online(tracks: self.index.size, transcode: info.ffmpeg)
                    self.checkedAt = Date()
                }
            } catch {
                let reason = error.localizedDescription
                await MainActor.run { self.status = .offline(reason: reason, cachedTracks: self.index.size) }
            }
        }
    }

    func find(_ track: Track) -> StemEntry? {
        let idx = index
        if let k = track.stemKey { return idx.entries.first { $0.key == k } }
        let hit = idx.find(track.partFile, durationMs: track.durationMs)
        let learned = idx.roots.joined(separator: "\n")
        if learned != settings.stemRoots { settings.stemRoots = learned }
        return hit
    }

    /// Format the server will actually send: Opus needs ffmpeg there.
    func effectiveQuality(_ q: StemQuality) -> StemQuality {
        q.fmt == "opus" && !transcode && status.isOnline ? .Lossless : q
    }

    func fileUrl(_ entry: StemEntry, _ stem: StemFile, _ quality: StemQuality) -> URL? {
        let path = (entry.key.split(separator: "/", omittingEmptySubsequences: false).map(String.init) + [stem.file])
            .map(urlEncode).joined(separator: "/")
        guard let u = URL(string: base + "/file/" + path) else { return nil }
        var params: [(String, Any?)] = []
        if quality.fmt == "opus" { params += [("fmt", "opus"), ("br", quality.kbps)] }
        if !settings.stemToken.isEmpty { params.append(("token", settings.stemToken)) }
        return u.with(params)
    }

    /// Asks the server to start the Opus transcodes now, so the stems are ready when we need them.
    func prepare(_ entry: StemEntry, _ quality: StemQuality) {
        guard quality.fmt == "opus" else { return }
        Task.detached { [self] in
            guard let u = try? url("/api/prepare", [("key", entry.key), ("fmt", "opus"), ("br", quality.kbps)]) else { return }
            var r = req(u, method: "POST")
            r.httpBody = Data()
            _ = try? await http.session.data(for: r)
        }
    }

    /// "Stem Now": asks stemd (through stemserve) to separate these songs before the rest of the
    /// library. Returns how many it queued: the ones with no stems yet.
    func stemNow(_ tracks: [Track]) async throws -> Int {
        var seen = Set<String>()
        let paths = tracks.filter { find($0) == nil }.compactMap(\.partFile).filter { seen.insert($0).inserted }
        if paths.isEmpty { return 0 }
        var r = req(try url("/api/stem"), method: "POST")
        r.setValue("application/json", forHTTPHeaderField: "Content-Type")
        r.httpBody = try JSONSerialization.data(withJSONObject: ["paths": paths])
        let (body, resp) = try await http.session.data(for: r)
        let code = (resp as? HTTPURLResponse)?.statusCode ?? 0
        if !(200..<300).contains(code) {
            let msg = (try? JSONSerialization.jsonObject(with: body) as? [String: Any])?["error"] as? String
            throw HttpError(code: code, message: code == 404 ? "stemserve needs updating for Stem Now" : msg ?? "stemserve answered \(code)")
        }
        await MainActor.run { paths.forEach { self.pending.insert($0) } }
        Task.detached { [self] in await pollQueue() }
        return paths.count
    }

    /// Per-stem RMS envelopes (0..1 of full scale), [buckets] values each.
    func peaks(_ entry: StemEntry, buckets: Int) async -> [String: [Float]]? {
        guard let u = try? url("/api/peaks", [("key", entry.key), ("n", buckets)]),
              let body = try? await http.data(req(u)),
              let obj = try? JSONSerialization.jsonObject(with: body) as? [String: Any],
              let stems = obj["stems"] as? [String: [Any]] else { return nil }
        return stems.mapValues { arr in arr.map { Float(($0 as? NSNumber)?.intValue ?? 0) / 1000 } }
    }
}

/// A simple async mutex: one closure at a time, in arrival order.
actor AsyncLock {
    private var busy = false
    private var waiters: [CheckedContinuation<Void, Never>] = []

    func run<T>(_ body: () async throws -> T) async rethrows -> T {
        if busy {
            await withCheckedContinuation { waiters.append($0) }
        }
        busy = true
        defer {
            if waiters.isEmpty { busy = false } else { waiters.removeFirst().resume() }
        }
        return try await body()
    }
}