import AVFoundation
import Foundation

enum HermesConnectionState: Equatable {
    case disconnected
    case connecting
    case connected
    case failed(String)
}

enum HermesPlaybackState: Equatable {
    case idle
    case playing
    case failed(String)
}

@MainActor
final class AudioStreamManager: NSObject, ObservableObject {
    @Published private(set) var connectionState: HermesConnectionState = .disconnected
    @Published private(set) var playbackState: HermesPlaybackState = .idle

    private let endpoint: URL
    private let session: URLSession
    private var task: URLSessionWebSocketTask?
    private var avPlayer: AVAudioPlayer?

    private let rawPlaybackEngine = AVAudioEngine()
    private let rawPlaybackNode = AVAudioPlayerNode()
    private let rawPlaybackFormat = AVAudioFormat(
        commonFormat: .pcmFormatInt16,
        sampleRate: 16_000,
        channels: 1,
        interleaved: true
    )!

    init(endpoint: URL) {
        self.endpoint = endpoint
        self.session = URLSession(configuration: .default)
        super.init()
        configureRawPlaybackEngine()
    }

    func connect() {
        guard task == nil else { return }

        connectionState = .connecting
        let task = session.webSocketTask(with: endpoint)
        self.task = task
        task.resume()
        connectionState = .connected
        receiveLoop()
    }

    func disconnect() {
        task?.cancel(with: .goingAway, reason: nil)
        task = nil
        connectionState = .disconnected
    }

    func sendPCMChunk(_ data: Data) {
        guard let task else { return }

        task.send(.data(data)) { [weak self] error in
            guard let error else { return }
            Task { @MainActor in
                self?.connectionState = .failed(error.localizedDescription)
            }
        }
    }

    func sendText(_ text: String) {
        guard let task else { return }

        task.send(.string(text)) { [weak self] error in
            guard let error else { return }
            Task { @MainActor in
                self?.connectionState = .failed(error.localizedDescription)
            }
        }
    }

    /// Optional protocol hint for servers that distinguish press/release.
    func sendControlEvent(_ name: String) {
        sendText(#"{"type":"\#(name)"}"#)
    }

    private func receiveLoop() {
        task?.receive { [weak self] result in
            Task { @MainActor in
                guard let self else { return }

                switch result {
                case .success(let message):
                    self.handle(message)
                    self.receiveLoop()
                case .failure(let error):
                    self.task = nil
                    self.connectionState = .failed(error.localizedDescription)
                }
            }
        }
    }

    private func handle(_ message: URLSessionWebSocketTask.Message) {
        switch message {
        case .data(let data):
            if data.isWAV {
                playWAV(data)
            } else {
                playRawPCMChunk(data)
            }
        case .string(let text):
            handleControlMessage(text)
        @unknown default:
            break
        }
    }

    private func handleControlMessage(_ text: String) {
        if text == "done" || text.contains(#""type":"done""#) {
            playbackState = .idle
        }
    }

    private func playWAV(_ data: Data) {
        do {
            try configurePlaybackSession()
            avPlayer = try AVAudioPlayer(data: data)
            avPlayer?.delegate = self
            avPlayer?.prepareToPlay()
            avPlayer?.play()
            playbackState = .playing
        } catch {
            playbackState = .failed(error.localizedDescription)
        }
    }

    private func configureRawPlaybackEngine() {
        rawPlaybackEngine.attach(rawPlaybackNode)
        rawPlaybackEngine.connect(rawPlaybackNode, to: rawPlaybackEngine.mainMixerNode, format: rawPlaybackFormat)
    }

    private func playRawPCMChunk(_ data: Data) {
        do {
            try configurePlaybackSession()

            if !rawPlaybackEngine.isRunning {
                try rawPlaybackEngine.start()
            }
            if !rawPlaybackNode.isPlaying {
                rawPlaybackNode.play()
            }

            guard let buffer = data.makePCMBuffer(format: rawPlaybackFormat) else { return }
            rawPlaybackNode.scheduleBuffer(buffer, completionHandler: nil)
            playbackState = .playing
        } catch {
            playbackState = .failed(error.localizedDescription)
        }
    }

    private func configurePlaybackSession() throws {
        let session = AVAudioSession.sharedInstance()
        try session.setCategory(.playAndRecord, mode: .default, options: [.defaultToSpeaker])
        try session.setActive(true)
    }
}

extension AudioStreamManager: AVAudioPlayerDelegate {
    nonisolated func audioPlayerDidFinishPlaying(_ player: AVAudioPlayer, successfully flag: Bool) {
        Task { @MainActor in
            self.playbackState = .idle
        }
    }
}

private extension Data {
    var isWAV: Bool {
        count >= 12 &&
            self[0] == 0x52 && self[1] == 0x49 && self[2] == 0x46 && self[3] == 0x46 &&
            self[8] == 0x57 && self[9] == 0x41 && self[10] == 0x56 && self[11] == 0x45
    }

    func makePCMBuffer(format: AVAudioFormat) -> AVAudioPCMBuffer? {
        let bytesPerFrame = Int(format.streamDescription.pointee.mBytesPerFrame)
        guard bytesPerFrame > 0 else { return nil }

        let frameCount = AVAudioFrameCount(count / bytesPerFrame)
        guard let buffer = AVAudioPCMBuffer(pcmFormat: format, frameCapacity: frameCount) else { return nil }
        buffer.frameLength = frameCount

        let audioBuffer = buffer.audioBufferList.pointee.mBuffers
        guard let destination = audioBuffer.mData else { return nil }

        withUnsafeBytes { rawBuffer in
            guard let source = rawBuffer.baseAddress else { return }
            destination.copyMemory(from: source, byteCount: Int(audioBuffer.mDataByteSize))
        }

        return buffer
    }
}