import Metal
import MetalKit
import CoreGraphics
import AVFoundation

public class RenderableVideoHelper {
    
    public struct FrameInput {
        public let width: Int
        public let height: Int
        public let frameIndex: Int
        public let framesPerSecond: TimeInterval
        public let deltaTime: TimeInterval
        public let timeElapsed: TimeInterval
        public let totalFrameCount: Int
        public let completionRatio: Float
        
        public init(
            width: Int,
            height: Int,
            frameIndex: Int,
            framesPerSecond: TimeInterval,
            deltaTime: TimeInterval,
            timeElapsed: TimeInterval,
            totalFrameCount: Int,
            completionRatio: Float) {
                
            self.width = width
            self.height = height
            self.frameIndex = frameIndex
            self.framesPerSecond = framesPerSecond
            self.deltaTime = deltaTime
            self.timeElapsed = timeElapsed
            self.totalFrameCount = totalFrameCount
            self.completionRatio = completionRatio
        }
    }
    
    private var renderables: [Renderable] = []
    
    // These determine the output values
    private var outputURL: URL?
    private var width: Int = 1024
    private var height: Int = 1024
    private var frameCount: Int = 60
    private var framesPerSecond: TimeInterval = 60
    
    private var offlineRenderHelper: OfflineRenderHelper?
    
    private var size: CGSize {
        return .init(width: CGFloat(width), height: CGFloat(height))
    }
    
    // Asset writer
    private var assetWriter: AVAssetWriter?
    private var assetWriterInput: AVAssetWriterInput?
    private var pixelBufferAdaptor: AVAssetWriterInputPixelBufferAdaptor?
    
    public init() { }
    
    public func configure(
        _ url: URL? = nil,
        avOutputSettings: [String: Any]? = nil,
        width: Int,
        height: Int,
        frameCount: Int,
        framesPerSecond: TimeInterval,
        renderables: [Renderable]) throws {
            
        self.width = width
        self.height = height
        self.frameCount = frameCount
        self.framesPerSecond = framesPerSecond
        self.outputURL = url ?? Adamantium.createRandomFileURLInDocuments("mp4")
        self.offlineRenderHelper = OfflineRenderHelper(width: width, height: height)
        self.renderables = renderables
            
        var avDefaultOutputSettings: [String: Any] = [
            AVVideoCodecKey: AVVideoCodecType.h264,
            AVVideoWidthKey: size.width,
            AVVideoHeightKey: size.height,
            AVVideoCompressionPropertiesKey: [
                AVVideoAverageBitRateKey: 10_000_000,               // Set a higher bitrate for better quality
                AVVideoMaxKeyFrameIntervalKey: framesPerSecond,     // Adjust for smoother playback
                AVVideoQualityKey: 0.5                              // Maximum quality (scale is 0.0 to 1.0)
            ]
        ]
            
        // No matter what the width and height will be determined by the configure parameters
        var outputSettings = avOutputSettings ?? avDefaultOutputSettings
        outputSettings[AVVideoWidthKey] = size.width
        outputSettings[AVVideoHeightKey] = size.height
        
        let sourcePixelBufferAttributesDictionary: [String: Any] = [
            kCVPixelBufferPixelFormatTypeKey as String: kCVPixelFormatType_32BGRA,
            kCVPixelBufferWidthKey as String: size.width,
            kCVPixelBufferHeightKey as String: size.height
        ]
            
        guard let outputURL else { return }
        
        do {
            self.assetWriter = try AVAssetWriter(outputURL: outputURL, fileType: .mp4)
            self.assetWriterInput = AVAssetWriterInput(mediaType: .video, outputSettings: outputSettings)
            guard let assetWriter, let assetWriterInput else { return }
            if assetWriter.canApply(outputSettings: outputSettings, forMediaType: .video), assetWriter.canAdd(assetWriterInput) {
                assetWriter.add(assetWriterInput)
            } else {
                throw RenderError.failedInitAssetWriter
            }
            
            self.pixelBufferAdaptor = AVAssetWriterInputPixelBufferAdaptor(
                assetWriterInput: assetWriterInput,
                sourcePixelBufferAttributes: sourcePixelBufferAttributesDictionary
            )
        } catch {
            throw RenderError.failedInitAssetWriter
        }
    }
    
    public func render(update: @escaping (FrameInput) -> Void) async throws -> URL {
        
        try await withCheckedThrowingContinuation { continuation in
            var hasUsedContinuation: Bool = false
            
            func throwError(_ error: RenderError) {
                guard !hasUsedContinuation else { return }
                hasUsedContinuation = true
                return continuation.resume(throwing: error)
            }
            
            guard
                let outputURL,
                let assetWriter,
                let assetWriterInput,
                let pixelBufferAdaptor,
                assetWriter.startWriting(),
                let pixelBufferPool = pixelBufferAdaptor.pixelBufferPool
            else {
                throwError(RenderError.failedWrite)
                return
            }
            
            assetWriter.startSession(atSourceTime: .zero)
            
            var frameIndex: Int = 0
            let size = size
            let width = width
            let height = height
            let frameRate = framesPerSecond
            let totalFrames = frameCount
            let renderables = renderables
            
            let context = CIContext()
            let queue = DispatchQueue(label: "mediaInputQueue")
            guard let pixelBuffer = try? ImageUtility.createPixelPuffer(size: size) else {
                throwError(RenderError.failedCreatePixelBuffer)
                return
            }
            
            func endWriting() {
                assetWriterInput.markAsFinished()
                assetWriter.finishWriting {
                    guard !hasUsedContinuation else { return }
                    hasUsedContinuation = true
                    if assetWriter.status == .completed {
                        continuation.resume(returning: outputURL)
                    } else {
                        continuation.resume(throwing: RenderError.failedWrite)
                    }
                }
            }
            
            assetWriterInput.requestMediaDataWhenReady(on: queue) { [weak self] in
                guard
                    let self,
                    let stillRenderHelper = offlineRenderHelper,
                    let assetWriterInput = self.assetWriterInput
                else { return }
                
                while assetWriterInput.isReadyForMoreMediaData {
                    /**
                     autorelease will get rid of any frame resources created during
                     this loop iteration.
                     
                     Without this the frame resources will pile up in memeory
                     which will eventually cause a crash in the consumer app
                     on anything over like 1 second
                    
                     Now with autorelease pool the memory usage of creating
                     a video hovers around 100MB on average, with one renderable.
                    */
                    autoreleasepool {
                        /**
                         This is where updates to state
                         for the current frame actually happen
                        */
                        let deltaTime: TimeInterval = 1.0 / frameRate
                        let timeElapsed = TimeInterval(frameIndex) * deltaTime
                        let completionRatio = Float(frameIndex + 1) / Float(totalFrames)
                        
                        let frameInput = FrameInput(
                            width: width,
                            height: height,
                            frameIndex: frameIndex,
                            framesPerSecond: frameRate,
                            deltaTime: deltaTime,
                            timeElapsed: timeElapsed,
                            totalFrameCount: totalFrames,
                            completionRatio: completionRatio
                        )
                        update(frameInput)
                        
                        /**
                         This is where those updates to state are
                         rendered into a metal texture.
                        */
                        stillRenderHelper.commitRenderableLayers(layers: renderables)
                        
                        guard
                            let texture = stillRenderHelper.lastRender,
                            let ciImage = try? StillCaptureUtility.ciImage(texture: texture)
                        else {
                            throwError(RenderError.failedToCreateCIImage)
                            return
                        }
                        
                        /**
                         This is where the rendered frame will be added to the target video
                        */
                        CVPixelBufferLockBaseAddress(pixelBuffer, [])
                        context.render(ciImage, to: pixelBuffer)
                        CVPixelBufferUnlockBaseAddress(pixelBuffer, [])
                        
                        let presentationTime = CMTime(value: Int64(frameIndex), timescale: Int32(frameRate))
                        pixelBufferAdaptor.append(pixelBuffer, withPresentationTime: presentationTime)
                        
                        if frameIndex < (totalFrames - 1) {
                            frameIndex += 1
                        } else {
                            endWriting()
                        }
                    }
                }
            }
        }
    }
    
    enum RenderError: LocalizedError {
        case failedWrite
        case failedInitAssetWriter
        case failedCreatePixelBuffer
        case failedToCreateCIImage
        
        var errorDescription: String? {
            switch self {
            case .failedWrite: "Failed to write media"
            case .failedInitAssetWriter: "Failed to init movie generator"
            case .failedCreatePixelBuffer: "Failed to create pixel buffer"
            case .failedToCreateCIImage: "Failed to create CIImage"
            }
        }
    }
}