import AVFoundation
import Foundation
import UIKit

struct ImageMovieGenerator {
    let image: UIImage
    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)

            let frameRate: Int32 = 1
            let totalFrames = Int32(targetDuration * Double(frameRate))

            let queue = DispatchQueue(label: "mediaInputQueue")
            assetWriterInput.requestMediaDataWhenReady(on: queue) {
                var frameCount: Int32 = 0
                while assetWriterInput.isReadyForMoreMediaData, frameCount < totalFrames {
                    let presentationTime = CMTimeMake(value: Int64(frameCount), timescale: Int32(frameRate))
                    guard let pixelBuffer = pixelBufferFromImage(image: image, pixelBufferPool: pixelBufferPool, size: size),
                          pixelBufferAdaptor.append(pixelBuffer, withPresentationTime: presentationTime)
                    else {
                        return continuation.resume(throwing: RenderError.failedCreatePixelBuffer)
                    }
                    frameCount += 1
                }

                assetWriterInput.markAsFinished()
                assetWriter.finishWriting {
                    if assetWriter.status == .completed {
                        continuation.resume(returning: outputURL)
                    } else {
                        continuation.resume(throwing: RenderError.failedWrite)
                    }
                }
            }
        }
    }

    init(image: UIImage, targetDuration: Double, outputURL: URL) throws {
        self.image = image
        self.targetDuration = targetDuration
        self.size = image.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"
            }
        }
    }

    func pixelBufferFromImage(image: UIImage, pixelBufferPool: CVPixelBufferPool, size: CGSize) -> CVPixelBuffer? {
        var pixelBufferOut: CVPixelBuffer?
        let status = CVPixelBufferPoolCreatePixelBuffer(
            kCFAllocatorDefault,
            pixelBufferPool,
            &pixelBufferOut
        )
        guard status == kCVReturnSuccess, let pixelBuffer = pixelBufferOut else {
            return nil
        }

        CVPixelBufferLockBaseAddress(pixelBuffer, [])
        let data = CVPixelBufferGetBaseAddress(pixelBuffer)
        let rgbColorSpace = CGColorSpaceCreateDeviceRGB()
        let context = CGContext(
            data: data,
            width: Int(size.width),
            height: Int(size.height),
            bitsPerComponent: 8,
            bytesPerRow: CVPixelBufferGetBytesPerRow(pixelBuffer),
            space: rgbColorSpace,
            bitmapInfo: CGImageAlphaInfo.noneSkipFirst.rawValue
        )

        guard let cgImage = image.cgImage else { return nil }
        let rectangle = CGRect(x: 0, y: 0, width: size.width, height: size.height)
        context?.draw(cgImage, in: rectangle)

        CVPixelBufferUnlockBaseAddress(pixelBuffer, [])

        return pixelBuffer
    }
}