import AVFAudio
import AVFoundation
import Foundation
import QuartzCore
import os

// MARK: - Recording State

enum AudioRecordingState: Sendable {
    case idle
    case recording(duration: TimeInterval, micLevel: Float)
    case paused(url: URL, duration: TimeInterval, waveform: [Float])
    case finished(url: URL, duration: TimeInterval, waveform: [Float])
}

// MARK: - AudioRecorder

/// Records voice messages using AVAudioRecorder.
/// Exposes mic level for animation and collects waveform samples.
@MainActor
final class AudioRecorder: NSObject {

    private let logger = Logger(subsystem: "com.rosetta.messenger", category: "AudioRecorder")

    private(set) var state: AudioRecordingState = .idle
    private(set) var micLevel: Float = 0

    var onLevelUpdate: ((TimeInterval, Float) -> Void)?
    var onFinished: ((URL, TimeInterval, [Float]) -> Void)?

    private var recorder: AVAudioRecorder?
    private var displayLink: CADisplayLink?
    private var waveformSamples: [Float] = []
    private var lastSampleTime: TimeInterval = 0
    private let sampleInterval: TimeInterval = 1.0 / 30.0

    // Multi-segment support for "record more" after preview
    private var previousSegments: [(url: URL, duration: TimeInterval)] = []
    private var accumulatedDuration: TimeInterval = 0

    private static let recordingSettings: [String: Any] = [
        AVFormatIDKey: Int(kAudioFormatMPEG4AAC),
        AVSampleRateKey: 48000,
        AVNumberOfChannelsKey: 1,
        AVEncoderAudioQualityKey: AVAudioQuality.high.rawValue,
        AVEncoderBitRateKey: 64000
    ]

    private var fileURL: URL {
        let tmp = FileManager.default.temporaryDirectory
        return tmp.appendingPathComponent("rosetta_voice_\(UUID().uuidString).m4a")
    }

    @discardableResult
    func startRecording() -> Bool {
        guard case .idle = state else {
            logger.warning("[AudioRecorder] startRecording called while not idle")
            return false
        }

        let session = AVAudioSession.sharedInstance()
        do {
            try session.setCategory(.playAndRecord, mode: .default, options: [.defaultToSpeaker])
            try session.setActive(true)
        } catch {
            logger.error("[AudioRecorder] Audio session failed: \(error.localizedDescription)")
            return false
        }

        let url = fileURL
        do {
            let rec = try AVAudioRecorder(url: url, settings: Self.recordingSettings)
            rec.isMeteringEnabled = true
            rec.delegate = self
            rec.prepareToRecord()
            guard rec.record() else {
                logger.error("[AudioRecorder] record() returned false")
                return false
            }
            recorder = rec
            waveformSamples = []
            lastSampleTime = 0
            micLevel = 0
            previousSegments = []
            accumulatedDuration = 0
            state = .recording(duration: 0, micLevel: 0)
            startDisplayLink()
            logger.info("[AudioRecorder] Started: \(url.lastPathComponent)")
            return true
        } catch {
            logger.error("[AudioRecorder] Init failed: \(error.localizedDescription)")
            return false
        }
    }

    func stopRecording() {
        if let rec = recorder {
            let currentDuration = rec.currentTime
            rec.stop()
            stopDisplayLink()
            let currentURL = rec.url
            recorder = nil
            finishWithFile(url: currentURL, fileDuration: currentDuration)
        } else if case .paused(let url, let duration, _) = state, !previousSegments.isEmpty {
            // Recorder already stopped (preview after record more) — concatenate segments
            finishWithFile(url: url, fileDuration: duration)
        } else if case .paused(let url, let duration, _) = state {
            // Single segment, recorder stopped — just finish
            state = .finished(url: url, duration: duration, waveform: waveformSamples)
            onFinished?(url, duration, waveformSamples)
            logger.info("[AudioRecorder] Stopped (from preview): \(String(format: "%.1f", duration))s")
        }
    }

    private func finishWithFile(url: URL, fileDuration: TimeInterval) {
        if previousSegments.isEmpty {
            state = .finished(url: url, duration: fileDuration, waveform: waveformSamples)
            onFinished?(url, fileDuration, waveformSamples)
            logger.info("[AudioRecorder] Stopped: \(String(format: "%.1f", fileDuration))s")
        } else {
            var allURLs = previousSegments.map(\.url)
            allURLs.append(url)
            let totalDuration = accumulatedDuration + fileDuration

            concatenateAudioFiles(urls: allURLs) { [weak self] resultURL in
                guard let self else { return }
                if let resultURL {
                    self.state = .finished(url: resultURL, duration: totalDuration, waveform: self.waveformSamples)
                    self.onFinished?(resultURL, totalDuration, self.waveformSamples)
                    self.logger.info("[AudioRecorder] Stopped (merged): \(String(format: "%.1f", totalDuration))s")
                } else {
                    self.state = .finished(url: url, duration: fileDuration, waveform: self.waveformSamples)
                    self.onFinished?(url, fileDuration, self.waveformSamples)
                    self.logger.warning("[AudioRecorder] Merge failed, using last segment")
                }
                for segment in self.previousSegments {
                    try? FileManager.default.removeItem(at: segment.url)
                }
                self.previousSegments = []
                self.accumulatedDuration = 0
            }
        }
    }

    /// Stops recording and finalizes the M4A file for preview playback.
    /// The file is fully written (moov atom included) so AVAudioPlayer can load it.
    /// For single-segment: returns immediately with the file URL.
    /// For multi-segment (after "record more"): merges all segments and calls completion with merged URL.
    func pauseRecordingForPreview(completion: @escaping ((url: URL, duration: TimeInterval, waveform: [Float])?) -> Void) {
        guard let rec = recorder, rec.isRecording else {
            completion(nil)
            return
        }
        let currentDuration = rec.currentTime
        let currentURL = rec.url
        rec.stop()  // Finalize M4A — writes moov atom so file is playable
        stopDisplayLink()
        recorder = nil

        if previousSegments.isEmpty {
            // Single segment — return immediately
            let totalDuration = currentDuration
            state = .paused(url: currentURL, duration: totalDuration, waveform: waveformSamples)
            logger.info("[AudioRecorder] Paused for preview: \(String(format: "%.1f", totalDuration))s")
            completion((url: currentURL, duration: totalDuration, waveform: waveformSamples))
        } else {
            // Multi-segment — merge all segments + current into one file for preview
            var allURLs = previousSegments.map(\.url)
            allURLs.append(currentURL)
            let totalDuration = accumulatedDuration + currentDuration
            let waveform = waveformSamples

            concatenateAudioFiles(urls: allURLs) { [weak self] mergedURL in
                guard let self else { return }
                if let mergedURL {
                    // Clean up individual segment files
                    for segment in self.previousSegments {
                        try? FileManager.default.removeItem(at: segment.url)
                    }
                    try? FileManager.default.removeItem(at: currentURL)
                    self.previousSegments = []
                    self.accumulatedDuration = 0
                    self.pendingTrimRange = nil
                    self.state = .paused(url: mergedURL, duration: totalDuration, waveform: waveform)
                    self.logger.info("[AudioRecorder] Paused for preview (merged): \(String(format: "%.1f", totalDuration))s")
                    completion((url: mergedURL, duration: totalDuration, waveform: waveform))
                } else {
                    // Merge failed — show only current segment
                    self.state = .paused(url: currentURL, duration: currentDuration, waveform: waveform)
                    self.logger.warning("[AudioRecorder] Merge for preview failed, using last segment")
                    completion((url: currentURL, duration: currentDuration, waveform: waveform))
                }
            }
        }
    }

    /// Synchronous version for backward compatibility (single-segment only).
    @discardableResult
    func pauseRecordingForPreview() -> (url: URL, duration: TimeInterval, waveform: [Float])? {
        guard let rec = recorder, rec.isRecording else { return nil }
        guard previousSegments.isEmpty else {
            // Multi-segment requires async version
            logger.warning("[AudioRecorder] pauseRecordingForPreview sync called with segments — use async version")
            return nil
        }
        let duration = rec.currentTime
        let url = rec.url
        rec.stop()
        stopDisplayLink()
        let snapshot = (url: url, duration: duration, waveform: waveformSamples)
        state = .paused(url: snapshot.url, duration: snapshot.duration, waveform: snapshot.waveform)
        recorder = nil
        logger.info("[AudioRecorder] Paused for preview: \(String(format: "%.1f", duration))s")
        return snapshot
    }

    /// Whether there are previous segments from "record more" that need merging.
    var hasPreviousSegments: Bool { !previousSegments.isEmpty }

    /// Starts a new recording segment for "record more" after preview.
    /// If `trimRange` is provided, the previous segment is trimmed before saving.
    /// The new recording will be concatenated with the trimmed segment on send.
    @discardableResult
    func resumeRecording(trimRange: ClosedRange<TimeInterval>? = nil) -> Bool {
        guard case .paused(let prevURL, let prevDuration, _) = state else { return false }

        // If trimmed, export only the trim range; otherwise keep full file
        if let trim = trimRange, trim.lowerBound > 0.1 || trim.upperBound < prevDuration - 0.1 {
            let trimmedDuration = trim.upperBound - trim.lowerBound
            previousSegments.append((url: prevURL, duration: trimmedDuration))
            pendingTrimRange = trim
        } else {
            previousSegments.append((url: prevURL, duration: prevDuration))
            pendingTrimRange = nil
        }
        accumulatedDuration = previousSegments.reduce(0) { $0 + $1.duration }

        let session = AVAudioSession.sharedInstance()
        do {
            try session.setCategory(.playAndRecord, mode: .default, options: [.defaultToSpeaker])
            try session.setActive(true)
        } catch {
            logger.error("[AudioRecorder] Resume session failed: \(error)")
            return false
        }

        let url = fileURL
        do {
            let rec = try AVAudioRecorder(url: url, settings: Self.recordingSettings)
            rec.isMeteringEnabled = true
            rec.delegate = self
            rec.prepareToRecord()
            guard rec.record() else {
                logger.error("[AudioRecorder] Resume record() failed")
                return false
            }
            recorder = rec
            state = .recording(duration: accumulatedDuration, micLevel: micLevel)
            startDisplayLink()
            logger.info("[AudioRecorder] Resumed recording (segment \(self.previousSegments.count + 1), trimmed: \(self.pendingTrimRange != nil))")
            return true
        } catch {
            logger.error("[AudioRecorder] Resume init failed: \(error)")
            return false
        }
    }

    /// Trim range to apply to the first segment during concatenation
    private var pendingTrimRange: ClosedRange<TimeInterval>?

    func currentRecordingSnapshot() -> (url: URL, duration: TimeInterval, waveform: [Float])? {
        if let rec = recorder {
            return (url: rec.url, duration: accumulatedDuration + rec.currentTime, waveform: waveformSamples)
        }
        switch state {
        case .paused(let url, let duration, let waveform):
            return (url: url, duration: duration, waveform: waveform)
        case .finished(let url, let duration, let waveform):
            return (url: url, duration: duration, waveform: waveform)
        default:
            return nil
        }
    }

    func cancelRecording() {
        if let rec = recorder {
            let url = rec.url
            rec.stop()
            try? FileManager.default.removeItem(at: url)
        }
        stopDisplayLink()
        // Clean up previous segment files
        for segment in previousSegments {
            try? FileManager.default.removeItem(at: segment.url)
        }
        logger.info("[AudioRecorder] Cancelled")
        recorder = nil
        reset()
    }

    func reset() {
        stopDisplayLink()
        recorder = nil
        micLevel = 0
        waveformSamples = []
        previousSegments = []
        accumulatedDuration = 0
        pendingTrimRange = nil
        state = .idle
    }

    func recordedData() -> Data? {
        guard case .finished(let url, _, _) = state else { return nil }
        return try? Data(contentsOf: url)
    }

    // MARK: - Display Link

    private func startDisplayLink() {
        let link = CADisplayLink(target: self, selector: #selector(displayLinkTick))
        link.add(to: .main, forMode: .common)
        displayLink = link
    }

    private func stopDisplayLink() {
        displayLink?.invalidate()
        displayLink = nil
    }

    @objc private func displayLinkTick() {
        guard let rec = recorder, rec.isRecording else { return }
        rec.updateMeters()
        let power = rec.averagePower(forChannel: 0)
        let normalized = Self.normalizeMicLevel(power)
        micLevel = normalized
        let duration = accumulatedDuration + rec.currentTime
        state = .recording(duration: duration, micLevel: normalized)
        if rec.currentTime - lastSampleTime >= sampleInterval {
            waveformSamples.append(normalized)
            lastSampleTime = rec.currentTime
        }
        onLevelUpdate?(duration, normalized)
    }

    static func normalizeMicLevel(_ power: Float) -> Float {
        let minDb: Float = -60
        let clamped = max(minDb, min(power, 0))
        let normalized = (clamped - minDb) / (-minDb)
        return normalized * normalized
    }

    static func requestMicrophonePermission() async -> Bool {
        let status = AVAudioSession.sharedInstance().recordPermission
        switch status {
        case .granted: return true
        case .undetermined:
            return await withCheckedContinuation { continuation in
                AVAudioSession.sharedInstance().requestRecordPermission { granted in
                    continuation.resume(returning: granted)
                }
            }
        case .denied: return false
        @unknown default: return false
        }
    }

    // MARK: - Audio Concatenation

    private func concatenateAudioFiles(urls: [URL], completion: @escaping @MainActor (URL?) -> Void) {
        let trimRange = pendingTrimRange
        let composition = AVMutableComposition()
        guard let track = composition.addMutableTrack(
            withMediaType: .audio,
            preferredTrackID: kCMPersistentTrackID_Invalid
        ) else {
            completion(nil)
            return
        }

        var insertTime = CMTime.zero
        for (index, url) in urls.enumerated() {
            let asset = AVURLAsset(url: url)
            guard let audioTrack = asset.tracks(withMediaType: .audio).first else { continue }
            do {
                // Apply trim range to the first segment (the original recording)
                let sourceRange: CMTimeRange
                if index == 0, let trim = trimRange {
                    let start = CMTime(seconds: trim.lowerBound, preferredTimescale: 44100)
                    let end = CMTime(seconds: trim.upperBound, preferredTimescale: 44100)
                    sourceRange = CMTimeRange(start: start, end: end)
                } else {
                    sourceRange = CMTimeRange(start: .zero, duration: asset.duration)
                }
                try track.insertTimeRange(sourceRange, of: audioTrack, at: insertTime)
                insertTime = CMTimeAdd(insertTime, sourceRange.duration)
            } catch {
                logger.error("[AudioRecorder] Concat insert failed: \(error)")
            }
        }

        let outputURL = FileManager.default.temporaryDirectory
            .appendingPathComponent("rosetta_voice_merged_\(UUID().uuidString).m4a")

        guard let exporter = AVAssetExportSession(asset: composition, presetName: AVAssetExportPresetAppleM4A) else {
            completion(nil)
            return
        }
        exporter.outputURL = outputURL
        exporter.outputFileType = .m4a
        exporter.exportAsynchronously {
            Task { @MainActor in
                completion(exporter.status == .completed ? outputURL : nil)
            }
        }
    }
}

extension AudioRecorder: AVAudioRecorderDelegate {
    nonisolated func audioRecorderDidFinishRecording(_ recorder: AVAudioRecorder, successfully flag: Bool) {
        Task { @MainActor in
            if !flag { logger.warning("[AudioRecorder] Finished unsuccessfully"); reset() }
        }
    }
    nonisolated func audioRecorderEncodeErrorDidOccur(_ recorder: AVAudioRecorder, error: Error?) {
        Task { @MainActor in
            logger.error("[AudioRecorder] Encode error: \(error?.localizedDescription ?? "unknown")")
            cancelRecording()
        }
    }
}
