import Foundation
import AVFoundation
import CoreMedia
import WebRTC

/// 远端视频录制器：把主控端收到的远程画面编码为 MP4 保存到 App 沙盒 Documents/WebRTCRecordings。
///
/// 与 Android 端 `VideoRecorder` 行为对齐，支持两种串流模式：
/// - WebRTC 模式：实现 `RTCVideoRenderer`，由 `WebRTCClient` 把远端视频轨 fan-out 一份过来，
///   直接拿到解码后的 `CVPixelBuffer`，用 `AVAssetWriter` + H.264 硬编码封装为 MP4。
/// - 自编码模式：由 `SelfCodecDecoder` 在硬解出 `CMSampleBuffer`（AVCC）后调用
///   `appendDecodedSampleBuffer(_:isKey:)`，复用已解码的 H.264 帧写盘。
///
/// 编码与写盘在独立串行队列执行，避免阻塞 WebRTC 网络/解码线程。
final class VideoRecorder: NSObject {

    enum RecorderError: Error {
        case alreadyRecording
        case notRecording
        case noTrackYet
        case writerInitFailed(String)
    }

    enum RecorderState {
        case idle
        case recording
    }

    /// 录制状态/结果回调（均回到主线程）
    var onStateChange: ((RecorderState) -> Void)?
    /// 停止并保存成功，savedPath 为完整文件路径
    var onSaved: ((String) -> Void)?
    /// 停止时未写入任何画面（例如未收到帧即停止）
    var onSavedEmpty: (() -> Void)?
    /// 真正错误（如 AVAssetWriter 初始化失败）
    var onError: ((String) -> Void)?

    private(set) var isRecording = false

    // 仅由录制队列访问
    private let recordQueue = DispatchQueue(label: "com.ttstd.video.recorder")
    private var assetWriter: AVAssetWriter?
    private var videoInput: AVAssetWriterInput?
    private var outputURL: URL?
    private var startedTime: CMTime?
    private var framesWritten = 0

    // MARK: - 公开控制

    /// 开始录制（WebRTC 模式需要 track 已经收到才能确定分辨率并开始写；
    /// 自编码模式通过 appendDecodedSampleBuffer 在收到首帧时自动起写）。
    /// 返回 true 表示已进入录制态；若当前正在录制返回 false。
    @discardableResult
    func start() -> Bool {
        if isRecording { return false }
        isRecording = true
        DispatchQueue.main.async { self.onStateChange?(.recording) }
        return true
    }

    /// 停止录制并保存。
    func stop() {
        guard isRecording else { return }
        isRecording = false
        recordQueue.async { [weak self] in
            self?.finalizeLocked()
        }
    }

    /// 释放（断开连接 / 页面销毁时调用），未保存内容会丢弃。
    func cancel() {
        guard isRecording else { return }
        isRecording = false
        recordQueue.async { [weak self] in
            self?.teardownLocked()
        }
    }

    // MARK: - WebRTC 模式入口（RTCVideoRenderer）

    /// 由 WebRTCClient 在收到远端视频轨后 add 此 renderer。
    /// 这里仅记录帧并投递到录制队列，真正的写盘在 ensureWriter 时启动。
    func handleWebRTCFrame(_ frame: RTCVideoFrame) {
        guard isRecording else { return }
        // buffer 为 RTCVideoFrameBuffer 协议，只有 RTCCVPixelBuffer 子类暴露 pixelBuffer
        guard let cvBuffer = frame.buffer as? RTCCVPixelBuffer else { return }
        let width = Int(cvBuffer.width)
        let height = Int(cvBuffer.height)
        guard width > 0, height > 0 else { return }
        let pts = CMTime(value: frame.timeStampNs, timescale: 1_000_000_000)

        recordQueue.async { [weak self] in
            guard let self else { return }
            self.appendPixelBuffer(cvBuffer.pixelBuffer, width: width, height: height, pts: pts)
        }
    }

    // MARK: - 自编码模式入口

    /// 自编码模式：SelfCodecDecoder 硬解出帧后调用。
    /// 首帧时依据 sampleBuffer 的分辨率启动写盘，AVAssetWriter 会重新编码为 H.264。
    func appendDecodedSampleBuffer(_ sample: CMSampleBuffer, isKey: Bool) {
        guard isRecording else { return }
        recordQueue.async { [weak self] in
            guard let self else { return }
            self.appendSampleBuffer(sample, isKey)
        }
    }

    // MARK: - 写入实现

    /// 把 CVPixelBuffer 包装成带时间戳的 CMSampleBuffer（WebRTC 模式写盘用）
    private func makeSampleBuffer(from pixelBuffer: CVPixelBuffer, pts: CMTime) -> CMSampleBuffer? {
        var formatDescription: CMVideoFormatDescription?
        guard CMVideoFormatDescriptionCreateForImageBuffer(
            allocator: kCFAllocatorDefault,
            imageBuffer: pixelBuffer,
            formatDescriptionOut: &formatDescription) == noErr,
              let formatDescription else { return nil }

        var sampleBuffer: CMSampleBuffer?
        var timing = CMSampleTimingInfo(
            duration: CMTime.invalid,
            presentationTimeStamp: pts,
            decodeTimeStamp: CMTime.invalid)
        let status = CMSampleBufferCreateForImageBuffer(
            allocator: kCFAllocatorDefault,
            imageBuffer: pixelBuffer,
            dataReady: true,
            makeDataReadyCallback: nil,
            refcon: nil,
            formatDescription: formatDescription,
            sampleTiming: &timing,
            sampleBufferOut: &sampleBuffer)
        guard status == noErr else { return nil }
        return sampleBuffer
    }

    private func appendPixelBuffer(_ pixelBuffer: CVPixelBuffer,
                                   width: Int, height: Int, pts: CMTime) {
        ensureWriterForPixelBuffer(width: width, height: height)
        guard let input = videoInput, input.isReadyForMoreMediaData else { return }
        var finalPts = pts
        if let start = startedTime {
            finalPts = CMTimeSubtract(pts, start)
        } else {
            startedTime = pts
            finalPts = .zero
        }
        guard let sample = makeSampleBuffer(from: pixelBuffer, pts: finalPts) else { return }
        let success = input.append(sample)
        if success { framesWritten += 1 }
    }

    private func appendSampleBuffer(_ sample: CMSampleBuffer, _ isKey: Bool = false) {
        ensureWriterForSampleBuffer(sample)
        guard let input = videoInput, input.isReadyForMoreMediaData else { return }
        if startedTime == nil {
            startedTime = CMSampleBufferGetPresentationTimeStamp(sample)
        }
        let pts = CMSampleBufferGetPresentationTimeStamp(sample)
        let finalPts = CMTimeSubtract(pts, startedTime ?? pts)
        // 用偏移后的时间戳构造新的 CMSampleBuffer 再写入
        guard let shifted = makeShiftedSampleBuffer(from: sample, pts: finalPts) else { return }
        let success = input.append(shifted)
        if success { framesWritten += 1 }
    }

    /// 复制 sampleBuffer 并把其 presentation/decode 时间戳替换为新的 pts（自编码模式写盘用）
    private func makeShiftedSampleBuffer(from sample: CMSampleBuffer, pts: CMTime) -> CMSampleBuffer? {
        var newSample: CMSampleBuffer?
        var timing = CMSampleTimingInfo(
            duration: CMTime.invalid,
            presentationTimeStamp: pts,
            decodeTimeStamp: CMTime.invalid)
        let status = CMSampleBufferCreateCopyWithNewTiming(
            allocator: kCFAllocatorDefault,
            sampleBuffer: sample,
            sampleTimingEntryCount: 1,
            sampleTimingArray: &timing,
            sampleBufferOut: &newSample)
        guard status == noErr else { return nil }
        return newSample
    }

    /// WebRTC 模式：用分辨率（NV12/ARGB 像素缓冲）初始化 AVAssetWriter
    private func ensureWriterForPixelBuffer(width: Int, height: Int) {
        guard assetWriter == nil else { return }
        guard let url = Self.buildOutputURL() else { return }
        outputURL = url
        do {
            let writer = try AVAssetWriter(outputURL: url, fileType: .mp4)
            let settings: [String: Any] = [
                AVVideoCodecKey: AVVideoCodecType.h264,
                AVVideoWidthKey: width,
                AVVideoHeightKey: height,
                AVVideoCompressionPropertiesKey: [
                    AVVideoAverageBitRateKey: min(max(width * height * 2, 1_000_000), 6_000_000),
                    AVVideoMaxKeyFrameIntervalKey: 30
                ]
            ]
            let input = AVAssetWriterInput(mediaType: .video, outputSettings: settings)
            input.expectsMediaDataInRealTime = true
            guard writer.canAdd(input) else {
                DispatchQueue.main.async {
                    self.onError?("AVAssetWriter 不支持该视频配置")
                }
                self.teardownLocked()
                return
            }
            writer.add(input)
            guard writer.startWriting() else {
                DispatchQueue.main.async {
                    self.onError?("AVAssetWriter 启动失败：\(writer.error?.localizedDescription ?? "未知错误")")
                }
                self.teardownLocked()
                return
            }
            writer.startSession(atSourceTime: .zero)
            assetWriter = writer
            videoInput = input
        } catch {
            DispatchQueue.main.async { self.onError?("录制初始化失败: \(error.localizedDescription)") }
            self.teardownLocked()
        }
    }

    /// 自编码模式：从解码后的 sampleBuffer 取分辨率（其格式描述），首帧时启动写盘。
    /// AVAssetWriterInput 配置为 H.264 压缩，会对解码后的帧重新编码并写盘，
    /// 无需手动解析 SPS/PPS。
    private func ensureWriterForSampleBuffer(_ sample: CMSampleBuffer) {
        if assetWriter != nil { return }
        guard let desc = CMSampleBufferGetFormatDescription(sample) else {
            NSLog("[VideoRecorder] 无法从 sampleBuffer 获取格式描述")
            return
        }
        let dims = CMVideoFormatDescriptionGetDimensions(desc)
        ensureWriterForSampleBuffer(width: Int(dims.width), height: Int(dims.height))
    }

    private func ensureWriterForSampleBuffer(width: Int, height: Int) {
        guard assetWriter == nil else { return }
        guard let url = Self.buildOutputURL() else { return }
        outputURL = url
        do {
            let writer = try AVAssetWriter(outputURL: url, fileType: .mp4)
            let settings: [String: Any] = [
                AVVideoCodecKey: AVVideoCodecType.h264,
                AVVideoWidthKey: width,
                AVVideoHeightKey: height,
                AVVideoCompressionPropertiesKey: [
                    AVVideoAverageBitRateKey: min(max(width * height * 2, 1_000_000), 6_000_000),
                    AVVideoMaxKeyFrameIntervalKey: 30
                ]
            ]
            let input = AVAssetWriterInput(mediaType: .video, outputSettings: settings)
            input.expectsMediaDataInRealTime = true
            guard writer.canAdd(input) else {
                DispatchQueue.main.async { self.onError?("AVAssetWriter 不支持该视频配置") }
                self.teardownLocked()
                return
            }
            writer.add(input)
            guard writer.startWriting() else {
                DispatchQueue.main.async {
                    self.onError?("AVAssetWriter 启动失败：\(writer.error?.localizedDescription ?? "未知错误")")
                }
                self.teardownLocked()
                return
            }
            writer.startSession(atSourceTime: .zero)
            assetWriter = writer
            videoInput = input
        } catch {
            DispatchQueue.main.async { self.onError?("录制初始化失败: \(error.localizedDescription)") }
            self.teardownLocked()
        }
    }

    private func finalizeLocked() {
        let written = framesWritten
        let url = outputURL
        videoInput?.markAsFinished()
        assetWriter?.finishWriting { [weak self] in
            guard let self else { return }
            self.teardownLocked()
            DispatchQueue.main.async {
                self.onStateChange?(.idle)
                if written > 0, let url {
                    self.onSaved?(url.path)
                } else {
                    self.onSavedEmpty?()
                }
            }
        }
    }

    private func teardownLocked() {
        assetWriter = nil
        videoInput = nil
        outputURL = nil
        startedTime = nil
        framesWritten = 0
    }

    // MARK: - 输出路径

    private static func buildOutputURL() -> URL? {
        guard let documents = FileManager.default.urls(for: .documentDirectory, in: .userDomainMask).first else {
            return nil
        }
        let dir = documents.appendingPathComponent("WebRTCRecordings", isDirectory: true)
        do {
            try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)
        } catch {
            NSLog("[VideoRecorder] 创建目录失败: \(error.localizedDescription)")
            return nil
        }
        let df = DateFormatter()
        df.dateFormat = "yyyyMMdd_HHmmss"
        let name = "rec_\(df.string(from: Date())).mp4"
        return dir.appendingPathComponent(name)
    }
}

// MARK: - RTCVideoRenderer（WebRTC 模式）

extension VideoRecorder: RTCVideoRenderer {
    func setSize(_ size: CGSize) {
        // 写盘分辨率在收到首帧 CVPixelBuffer 时确定，这里无需处理
    }

    func renderFrame(_ frame: RTCVideoFrame?) {
        guard let frame else { return }
        handleWebRTCFrame(frame)
    }
}
