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

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
    }
}