import Foundation
import AVFoundation
import AVKit
import CoreImage
import CoreImage.CIFilterBuiltins

public class VideoWatermarkUtility {
    
    public enum UtilityError: Error {
        case CouldNotCreateWatermarkingSession
        case WatermarkingSessionFailed
    }
    
    public struct WatermarkTransform {
        let opacity: CGFloat
        let widthRatioScale: CGSize
        let widthRatioTranslation: CGPoint
        
        public init(
            opacity: CGFloat,
            widthRatioScale: CGSize,
            widthRatioTranslation: CGPoint) {
                
            self.opacity = opacity
            self.widthRatioScale = widthRatioScale
            self.widthRatioTranslation = widthRatioTranslation
        }
    }
    
    // Main function to extract image, apply watermark, and create the watermarked video with audio
    public static func watermarkImageVideo(
        inputURL: URL?,
        markImage: CIImage?,
        markTransform: WatermarkTransform) async throws -> URL? {
        
        guard 
            let inputURL,
            let markImage,
            let videoInputs = try await getVideoAssetInputs(inputURL),
            let stillImage = VideoCompositionUtility.extractStillImageFromVideo(videoURL: inputURL),
            let outputURL = Adamantium.createRandomFileURLInDocuments("mp4")
        else { return nil }
            
        let waterMarkedImage = applyWatermarkToSingleImage(
            .init(cgImage: stillImage),
            size: .init(width: stillImage.width, height: stillImage.height),
            markTransform: markTransform,
            markImage: markImage)
            
        let duration = videoInputs.timerange.duration.seconds.rounded(.up)
            
        let movieGen = try ImageMovieGenerator(
            image: waterMarkedImage,
            targetDuration: duration,
            outputURL: outputURL)
            
        let renderedMovieURL = try await movieGen.render()
        let videoComposer = VideoCompositionUtility()
        return await videoComposer.combineAudioWithVideo(renderedMovieURL, inputURL)
    }
    
    public static func createWatermarkedAsset(
    inputURL: URL?,
    markImage: CIImage?,
    markTransform: WatermarkTransform,
    markFadeOutDuration: TimeInterval) async throws -> AVAssetExportSession? {
        
        guard
            let markImage,
            let inputs = try await getVideoAssetInputs(inputURL),
            let composition = createWatermarkingComposition(
                inputs, markTransform: markTransform, maskImage: markImage, markFadeOutDuration: markFadeOutDuration)
        else { return nil }
        
        return await exportVideoComposition(videoInputs: inputs, videoComposition: composition)
    }
    
    public static func createWatermarkedAssetForFirstNSeconds(
        inputURL: URL?,
        markImage: CIImage?,
        markTransform: WatermarkTransform,
        markDuration: CMTime,
        markFadeOutDuration: TimeInterval
    ) async throws -> AVAssetExportSession? {
        
        guard let inputURL else {
            print("Adamantium: Error with inputURL in VideoWatermarkUtility")
            return nil
        }
        
        // 1. Extract first n seconds of the video
        guard let firstClipSession = try? await extractClip(
            from: inputURL,
            duration: markDuration),
            firstClipSession.status == .completed,
            let firstClipURL = firstClipSession.outputURL
        else {
            print("Adamantium: Error extracting clip in VideoWatermarkUtility")
            return nil
        }
        
        // 2. Apply watermark to the extracted clip
        guard let watermarkedClipSession = try? await createWatermarkedAsset(
            inputURL: firstClipURL,
            markImage: markImage,
            markTransform: markTransform,
            markFadeOutDuration: markFadeOutDuration),
            watermarkedClipSession.status == .completed,
            let watermarkedClipURL = watermarkedClipSession.outputURL
        else {
            print("Adamantium: Error watermarking asset in VideoWatermarkUtility")
            return nil
        }
        
        // 3. Concatenate the watermarked clip with the remaining part of the video
        guard let concatenatedVideoSession = try? await concatenateWatermarkedClip(
            watermarkedClipURL: watermarkedClipURL,
            originalVideoURL: inputURL)
        else {
            print("Adamantium: Error concatenating video in VideoWatermarkUtility")
            return nil
        }
        
        return concatenatedVideoSession
    }
}

private extension VideoWatermarkUtility {
    
    struct VideoCompositionInput {
        let size: CGSize
        let asset: AVAsset
        let track: AVAssetTrack
        let timerange: CMTimeRange
        let mixComposition: AVMutableComposition
        let compositionTrack: AVMutableCompositionTrack
    }
    
    static func getVideoAssetInputs(_ url: URL?) async throws -> VideoCompositionInput? {
        let mixComposition = AVMutableComposition()
        guard let url = url else { return nil }
        let asset = AVURLAsset(url: url)
        guard let videoTrack = try? await asset.loadTracks(withMediaType: .video).first else {
            print("Adamantium: Failed to get video track in VideoWatermarkUtility")
            return nil
        }

        guard let timerange = try? await videoTrack.load(.timeRange),
            let compositionVideoTrack = mixComposition
            .addMutableTrack(
                withMediaType: .video,
                preferredTrackID: kCMPersistentTrackID_Invalid)
        else {
            print("Adamantium: Failed to add video track to composition in VideoWatermarkUtility")
            return nil
        }
        
        do {
            try compositionVideoTrack.insertTimeRange(timerange, of: videoTrack, at: .zero)
        } catch {
            print("Adamantium: Error inserting time range: \(error) in VideoWatermarkUtility")
            return nil
        }

        // get video size
        guard let size = try? await videoTrack.load(.naturalSize) else {
            print("Adamantium: Error getting video track size in VideoWatermarkUtility")
            return nil
        }
        
        return VideoCompositionInput(
            size: size,
            asset: asset,
            track: videoTrack,
            timerange: timerange,
            mixComposition: mixComposition,
            compositionTrack: compositionVideoTrack)
    }
    
    static func createWatermarkingComposition(
        _ inputs: VideoCompositionInput,
        markTransform: WatermarkTransform,
        maskImage: CIImage,
        markFadeOutDuration: TimeInterval) -> AVVideoComposition? {
            
        let composition = AVVideoComposition(asset: inputs.asset) { request in
            // Fade Out Opacity
            let fadeOutStartTime = inputs.timerange.duration.seconds - markFadeOutDuration
            let currentTime = request.compositionTime.seconds
            let targetOpacity: CGFloat
            
            if currentTime >= fadeOutStartTime {
                let normalizedFadeOutCompletionTime = (currentTime - fadeOutStartTime) / markFadeOutDuration
                targetOpacity = 1.0 - normalizedFadeOutCompletionTime
            } else {
                targetOpacity = 1.0
            }
            
            let opacity = targetOpacity * markTransform.opacity
            
            // Watermark Transform
            let source = request.sourceImage.clampedToExtent()
            
            // Resize and position the image if needed
            let width = request.renderSize.width
            let watermarkSize = CGSize(width: width * markTransform.widthRatioScale.width,
                                       height: width * markTransform.widthRatioScale.height)
            
            // Position watermark based on the video size and ratioTranslation
            let translation = CGAffineTransform(
                translationX: width * markTransform.widthRatioTranslation.x,
                y: width * markTransform.widthRatioTranslation.y)
            
            // Scale the watermark
            let scale = CGAffineTransform(scaleX: watermarkSize.width / maskImage.extent.width,
                                          y: watermarkSize.height / maskImage.extent.height)
            
            // Apply transformations: scale and translation
            var transformedImage = maskImage
                .transformed(by: scale, highQualityDownsample: true)
                .transformed(by: translation)
            
            // Apply opacity using a color matrix filter
            let alphaFilter = CIFilter.colorMatrix()
            alphaFilter.inputImage = transformedImage
            alphaFilter.aVector = CIVector(x: 0, y: 0, z: 0, w: opacity)
            
            // Apply opacity to the alpha channel
            if let alphaImage = alphaFilter.outputImage {
                transformedImage = alphaImage
            }
            
            // Apply compositing filter to overlay the image onto the video frame
            let maxFilter = CIFilter.maximumCompositing()
            maxFilter.setValue(transformedImage, forKey: kCIInputImageKey)
            maxFilter.setValue(source, forKey: kCIInputBackgroundImageKey)
            
            // Get the blended frame
            if let outputImage = maxFilter.outputImage {
                request.finish(with: outputImage.cropped(to: request.sourceImage.extent), context: nil)
            } else {
                request.finish(with: source, context: nil)
            }
        }
        
        return composition
    }
    
    static func exportVideoComposition(
        videoInputs: VideoCompositionInput?,
        videoComposition: AVVideoComposition?) async -> AVAssetExportSession? {
        
        guard 
            let videoInputs,
            let videoComposition,
            let exportSession = AVAssetExportSession(
            asset: videoInputs.asset,
            presetName: AVAssetExportPresetHEVC1920x1080)
        else {
            return nil
        }
            
        // Get the document directory to save the exported video
        let exportURL = Adamantium.createRandomFileURLInDocuments("mp4")
        exportSession.outputURL = exportURL
        exportSession.outputFileType = .mp4
        exportSession.videoComposition = videoComposition
        await exportSession.export()
        return exportSession
    }
        
    // Helper to extract a clip of n seconds
    static func extractClip(from url: URL, duration: CMTime) async throws -> AVAssetExportSession? {
        let asset = AVURLAsset(url: url)
        let mixComposition = AVMutableComposition()
        
        guard let videoTrack = try await asset.loadTracks(withMediaType: .video).first else { return nil }
        
        // Define the time range for the first n seconds
        let timeRange = CMTimeRangeMake(start: .zero, duration: duration)
        
        guard let compositionVideoTrack = mixComposition.addMutableTrack(
            withMediaType: .video,
            preferredTrackID: kCMPersistentTrackID_Invalid
        ) else { return nil }
        
        // Insert the first n seconds of the video
        try compositionVideoTrack.insertTimeRange(timeRange, of: videoTrack, at: .zero)
        
        // Export the first n seconds to a file
        return try await exportComposition(mixComposition)
    }

    // Concatenate the watermarked clip with the remaining original video
    static func concatenateWatermarkedClip(
        watermarkedClipURL: URL,
        originalVideoURL: URL
    ) async throws -> AVAssetExportSession? {
        
        let mixComposition = AVMutableComposition()
        
        // Add watermarked video (first n seconds)
        let watermarkAsset = AVURLAsset(url: watermarkedClipURL)
        guard let watermarkedAsset = try await watermarkAsset.loadTracks(withMediaType: .video).first,
              let watermarkDuration = try? await watermarkAsset.load(.duration),
              let videoTrack = mixComposition.addMutableTrack(
                withMediaType: .video,
                preferredTrackID: kCMPersistentTrackID_Invalid),
              let audioTrack = mixComposition.addMutableTrack(
                withMediaType: .audio,
                preferredTrackID: kCMPersistentTrackID_Invalid)
        else {
            print("Adamantium: Error on getting video track in concatenate in VideoWatermarkUtility")
            return nil
        }
        
        let nSecondsTimeRange = CMTimeRange(start: .zero, duration: watermarkDuration)
        do {
            try videoTrack.insertTimeRange(nSecondsTimeRange, of: watermarkedAsset, at: .zero)
        } catch {
            print("Adamantium: Error inserting time range in concatenate in VideoWatermarkUtility")
            return nil
        }
        
        // Add the remaining part of the original video
        let originalAsset = AVURLAsset(url: originalVideoURL)
        guard
            let originalVideoTrack = try? await originalAsset.loadTracks(withMediaType: .video).first,
            let originalAudioTrack = try? await originalAsset.loadTracks(withMediaType: .audio).first,
            let originalVideoTrackDuration = try? await originalAsset.load(.duration)
        else {
            print("Adamantium: Error getting the original asset track in VideoWatermarkUtility")
            return nil
        }
        
        let remainingDuration = originalVideoTrackDuration - watermarkDuration
        let remainingRange = CMTimeRange(
            start: watermarkDuration,
            duration: remainingDuration
        )
        
        do {
            try videoTrack.insertTimeRange(remainingRange, of: originalVideoTrack, at: watermarkDuration)
            try audioTrack.insertTimeRange(
                .init(start: .zero, duration: originalVideoTrackDuration),
                of: originalAudioTrack,
                at: .zero)
            
             // Export the final video after concatenating
            return try await exportComposition(mixComposition)
        } catch {
            print("Adamantium: Error on composition of concatenated video: \(error) in VideoWatermarkUtility")
            return nil
        }
    }
    
    // Helper to export a composition to a file
    static func exportComposition(_ composition: AVMutableComposition) async throws -> AVAssetExportSession? {
        let exportSession = AVAssetExportSession(asset: composition, presetName: AVAssetExportPresetHEVC1920x1080)
        let outputURL = Adamantium.createRandomFileURLInDocuments("mp4")
        exportSession?.outputURL = outputURL
        exportSession?.outputFileType = .mp4
        exportSession?.timeRange = CMTimeRangeMake(start: .zero, duration: composition.duration)
        
        await exportSession?.export()
        return exportSession
    }
    
    static func applyWatermarkToSingleImage(
        _ sourceImage: CIImage,
        size: CGSize,
        markTransform: WatermarkTransform,
        markImage: CIImage) -> CIImage {
            
        let opacity = markTransform.opacity
        // Watermark Transform
        
        // Resize and position the image if needed
        let width = size.width
        let watermarkSize = CGSize(
            width: width * markTransform.widthRatioScale.width,
            height: width * markTransform.widthRatioScale.height)
        
        // Position watermark based on the video size and ratioTranslation
        let translation = CGAffineTransform(
            translationX: width * markTransform.widthRatioTranslation.x,
            y: width * markTransform.widthRatioTranslation.y)
        
        // Scale the watermark
        let scale = CGAffineTransform(
            scaleX: watermarkSize.width / markImage.extent.width,
            y: watermarkSize.height / markImage.extent.height)
            
        // Apply transformations: scale and translation
        var transformedImage = markImage
            .transformed(by: scale, highQualityDownsample: true)
            .transformed(by: translation)
        
        // Apply opacity using a color matrix filter
        let alphaFilter = CIFilter.colorMatrix()
        alphaFilter.inputImage = transformedImage
        alphaFilter.aVector = CIVector(x: 0, y: 0, z: 0, w: opacity)
        
        // Apply opacity to the alpha channel
        if let alphaImage = alphaFilter.outputImage {
            transformedImage = alphaImage
        }
        
        // Apply compositing filter to overlay the image onto the video frame
        let maxFilter = CIFilter.maximumCompositing()
        maxFilter.setValue(transformedImage, forKey: kCIInputImageKey)
        maxFilter.setValue(sourceImage, forKey: kCIInputBackgroundImageKey)
        
        guard let outputImage = maxFilter.outputImage else { return sourceImage }
            
        return outputImage
    }

}