import AVFoundation
import Foundation
import MetalKit

public class VideoCaptureUtility {
    // MARK: - Properties

    private let writerQueue = DispatchQueue(label: "com.myApp.AssetWriterQueue",
                                            autoreleaseFrequency: .workItem)

    public private(set) var isRecording = false
    public private(set) var recordingStartTime = TimeInterval.zero

    private let assetWriter: AVAssetWriter
    private let assetWriterVideoInput: AVAssetWriterInput
    private let assetWriterPixelBufferInput: AVAssetWriterInputPixelBufferAdaptor

    public var captureStatus: AVAssetWriter.Status {
        assetWriter.status
    }

    public init?(outputURL url: URL, size: CGSize) {
        do {
            assetWriter = try AVAssetWriter(outputURL: url, fileType: .mp4)
        } catch {
            print("VideoCaptureUtility: Failed to initialize AVAssetWriter with error: \(error)")
            return nil
        }

        let outputSettings: [String: Any] = [
            AVVideoCodecKey: AVVideoCodecType.hevc,
            AVVideoWidthKey: size.width,
            AVVideoHeightKey: size.height,
        ]

        let videoInput = AVAssetWriterInput(mediaType: .video, outputSettings: outputSettings)
        videoInput.expectsMediaDataInRealTime = true

        guard assetWriter.canAdd(videoInput) else {
            print("VideoCaptureUtility: Cannot add video input to AVAssetWriter.")
            return nil
        }
        assetWriter.add(videoInput)
        assetWriterVideoInput = videoInput

        let sourcePixelBufferAttributes: [String: Any] = [
            kCVPixelBufferPixelFormatTypeKey as String: kCVPixelFormatType_32BGRA,
            kCVPixelBufferWidthKey as String: size.width,
            kCVPixelBufferHeightKey as String: size.height,
        ]

        assetWriterPixelBufferInput = AVAssetWriterInputPixelBufferAdaptor(
            assetWriterInput: videoInput,
            sourcePixelBufferAttributes: sourcePixelBufferAttributes
        )
    }

    // MARK: - Recording Control

    public func startRecording() {
        writerQueue.sync {
            guard assetWriter.status != .writing else { return }

            if assetWriter.startWriting() {
                assetWriter.startSession(atSourceTime: .zero)
                recordingStartTime = CACurrentMediaTime()
                isRecording = true
            } else {
                print("VideoCaptureUtility: Failed to start writing. Status: \(assetWriter.status)")
            }
        }
    }

    public func endRecording(_ completionHandler: @escaping (URL) -> Void) {
        writerQueue.sync {
            guard isRecording else { return }
            isRecording = false

            let outputURL = assetWriter.outputURL
            assetWriterVideoInput.markAsFinished()

            assetWriter.finishWriting {
                DispatchQueue.main.async {
                    completionHandler(outputURL)
                }
            }
        }
    }

    /// Writes a frame from the given `MTLTexture`. If the texture size is larger or smaller than
    /// the pixel buffer, this will simply copy the region that fits (i.e., crop).
    ///
    /// - Parameters:
    ///   - texture: The Metal texture to copy from.
    ///   - customTimeIncrease: If provided, uses this time offset; otherwise uses `CACurrentMediaTime()`.
    ///   - isBufferAllocationOnly: Skip appending the buffer to `AVAssetWriter` if `true`.
    public func writeFrame(
        forTexture texture: MTLTexture,
        customTimeIncrease: CFTimeInterval? = nil,
        isBufferAllocationOnly: Bool = false
    ) {
        writerQueue.sync {
            guard isRecording else {
                print("Tried to write but is not recording")
                return
            }
            guard assetWriter.status == .writing else {
                print("Asset writer status is not writing")
                return
            }
            guard assetWriterVideoInput.isReadyForMoreMediaData else {
                print("Asset writer not ready for more media")
                return
            }

            guard let pixelBuffer = createPixelBufferFromPool() else { return }

            CVPixelBufferLockBaseAddress(pixelBuffer, [])

            // Copy the region that fits to avoid out-of-bounds.
            let success = copyTextureBytes(texture, toPixelBuffer: pixelBuffer)

            if success, !isBufferAllocationOnly {
                let presentationTime = frameTime(using: customTimeIncrease)
                assetWriterPixelBufferInput.append(pixelBuffer, withPresentationTime: presentationTime)
            } else {
                if isBufferAllocationOnly {
                    print("Adamantium: Allocating video buffer but not recording [Expected]")
                } else {
                    print("Adamantium: Not successful in writing to video pixel buffer")
                }
            }

            CVPixelBufferUnlockBaseAddress(pixelBuffer, [])
        }
    }
}

// MARK: - Private Helpers

private extension VideoCaptureUtility {
    func createPixelBufferFromPool() -> CVPixelBuffer? {
        guard let pool = assetWriterPixelBufferInput.pixelBufferPool else {
            print("VideoCaptureUtility: No pixel buffer pool found. Unable to retrieve buffer.")
            return nil
        }

        var pixelBufferOut: CVPixelBuffer?
        let status = CVPixelBufferPoolCreatePixelBuffer(nil, pool, &pixelBufferOut)

        if status != kCVReturnSuccess {
            print("VideoCaptureUtility: Failed to create pixel buffer from pool (status: \(status)).")
            return nil
        }

        return pixelBufferOut
    }

    /// ** Pixel Buffer Guard**:
    /// - If texture is bigger than the pixel buffer, only copy the portion that fits (i.e., crop).
    /// - If it's smaller, only copy the part of the texture that exists (the rest of the buffer remains unchanged).
    func copyTextureBytes(_ texture: MTLTexture, toPixelBuffer pixelBuffer: CVPixelBuffer) -> Bool {
        guard let baseAddress = CVPixelBufferGetBaseAddress(pixelBuffer) else {
            print("VideoCaptureUtility: Pixel buffer has no base address.")
            return false
        }

        let pbWidth = CVPixelBufferGetWidth(pixelBuffer)
        let pbHeight = CVPixelBufferGetHeight(pixelBuffer)

        // Determine how much we can actually copy without exceeding bounds.
        let copyWidth = min(texture.width, pbWidth)
        let copyHeight = min(texture.height, pbHeight)

        // If there's nothing to copy, bail out.
        guard copyWidth > 0, copyHeight > 0 else {
            print("VideoCaptureUtility: No overlapping region to copy.")
            return false
        }

        let bytesPerRow = CVPixelBufferGetBytesPerRow(pixelBuffer)

        // We'll copy only from (0,0) to (copyWidth,copyHeight).
        let region = MTLRegionMake2D(0, 0, copyWidth, copyHeight)

        texture.getBytes(
            baseAddress,
            bytesPerRow: bytesPerRow,
            from: region,
            mipmapLevel: 0
        )

        return true
    }

    func frameTime(using customTimeIncrease: CFTimeInterval?) -> CMTime {
        let frameTime: CFTimeInterval
        if let time = customTimeIncrease {
            frameTime = time
        } else {
            frameTime = CACurrentMediaTime() - recordingStartTime
        }
        return CMTimeMakeWithSeconds(frameTime, preferredTimescale: 600)
    }
}