import Foundation
import UIKit
import AVFoundation

class ImageMovieGenerator {
    let image: CIImage
    let size: CGSize
    let targetDuration: Double
    let outputURL: URL
    
    var assetWriter: AVAssetWriter
    var assetWriterInput: AVAssetWriterInput
    var pixelBufferAdaptor: AVAssetWriterInputPixelBufferAdaptor
    
    func render() async throws -> URL {
        try await withCheckedThrowingContinuation { continuation in
            guard assetWriter.startWriting(),
                  let pixelBufferPool = pixelBufferAdaptor.pixelBufferPool
            else {
                return continuation.resume(throwing: RenderError.failedWrite)
            }
            
            assetWriter.startSession(atSourceTime: .zero)
            
            var frameIndex: Int32 = 0
            let frameRate: Int32 = 1
            let totalFrames = Int32(targetDuration * Double(frameRate))
            
            let queue = DispatchQueue(label: "mediaInputQueue")
            assetWriterInput.requestMediaDataWhenReady(on: queue) { [weak self] in
                guard let self else { return }
                
                while self.assetWriterInput.isReadyForMoreMediaData {
                    // Copy the next sample buffer from your source media.
                    guard let nextSampleBuffer = ImageUtility.ciImageToCVPixelBuffer(self.image, size: self.size) else {
                        // Mark the input as finished.
                        self.assetWriterInput.markAsFinished()
                        break
                    }
                    
                    let presentationTime = CMTimeMake(value: Int64(frameIndex), timescale: Int32(frameRate))
                    self.pixelBufferAdaptor.append(nextSampleBuffer, withPresentationTime: presentationTime)
                    
                    if frameIndex < (totalFrames - 1) {
                        frameIndex += 1
                    } else {
                        self.assetWriterInput.markAsFinished()
                        self.assetWriter.finishWriting {
                            if self.assetWriter.status == .completed {
                                continuation.resume(returning: self.outputURL)
                            } else {
                                continuation.resume(throwing: RenderError.failedWrite)
                            }
                        }
                    }
                }
            }
        }
    }
    
    init(image: CIImage, targetDuration: Double, outputURL: URL) throws {
        self.image = image
        self.targetDuration = targetDuration
        self.size = image.extent.size
        self.outputURL = outputURL
        
        var avOutputSettings: [String: Any] = [
            AVVideoCodecKey: AVVideoCodecType.h264,
            AVVideoWidthKey: size.width,
            AVVideoHeightKey: size.height
        ]
        
        if let videoSettings = AVOutputSettingsAssistant(preset: .hevc1920x1080)?.videoSettings {
            avOutputSettings = videoSettings
            // Use the same orientation as the image
            avOutputSettings[AVVideoWidthKey] = size.width
            avOutputSettings[AVVideoHeightKey] = size.height
        }
        
        let sourcePixelBufferAttributesDictionary: [String: Any] = [
            kCVPixelBufferPixelFormatTypeKey as String: kCVPixelFormatType_32ARGB,
            kCVPixelBufferWidthKey as String: size.width,
            kCVPixelBufferHeightKey as String: size.height
        ]
        
        do {
            self.assetWriter = try AVAssetWriter(outputURL: outputURL, fileType: AVFileType.mp4)
            self.assetWriterInput = AVAssetWriterInput(mediaType: .video, outputSettings: avOutputSettings)
            if assetWriter.canApply(outputSettings: avOutputSettings, forMediaType: .video), assetWriter.canAdd(assetWriterInput) {
                assetWriter.add(assetWriterInput)
            } else {
                throw RenderError.failedInitMovieGenerator
            }
            
            self.pixelBufferAdaptor = AVAssetWriterInputPixelBufferAdaptor(
                assetWriterInput: assetWriterInput,
                sourcePixelBufferAttributes: sourcePixelBufferAttributesDictionary
            )
        } catch {
            throw RenderError.failedInitMovieGenerator
        }
    }
    
    enum RenderError: LocalizedError {
        case failedInitMovieGenerator
        case failedWrite
        case failedCreatePixelBuffer
        
        var errorDescription: String? {
            switch self {
            case .failedInitMovieGenerator: "Failed to init movie generator"
            case .failedWrite: "Failed to write media"
            case .failedCreatePixelBuffer: "Failed to create pixel buffer"
            }
        }
    }
}