import AVFoundation
import CoreImage
import Foundation
import UIKit

class VideoCompositionUtility {
    public enum VideoCompositionError: Error {
        case compositionExportFailed
        case creationSessionFailed
        case couldNotAddAudioTrack
        case couldNotAddVideoTrack
    }

    // Extract the still image from the video
    static func extractStillImageFromVideo(videoURL: URL) -> CGImage? {
        let asset = AVAsset(url: videoURL)
        let imageGenerator = AVAssetImageGenerator(asset: asset)
        imageGenerator.appliesPreferredTrackTransform = true

        // Generate image at the 0th second (since it's a still image throughout)
        let time = CMTime(seconds: .zero, preferredTimescale: 600)

        do {
            let cgImage = try imageGenerator.copyCGImage(at: time, actualTime: nil)
            return cgImage
        } catch {
            print("Adamantium: Error extracting image from video extractStillImageFromVideo(videoURL: URL) in VideoCompositionError: \(error)")
            return nil
        }
    }

    func combineAudioWithVideo(_ videoAssetURL: URL, _ audioAssetURL: URL) async -> URL? {
        let videoAsset = AVAsset(url: videoAssetURL)
        let audioAsset = AVAsset(url: audioAssetURL)
        let mixComposition = AVMutableComposition()

        // Get the video track from the video asset
        guard let videoTrack = videoAsset.tracks(withMediaType: .video).first else { return nil }

        // Create a video composition track
        guard let videoCompositionTrack = mixComposition
            .addMutableTrack(withMediaType: .video, preferredTrackID: kCMPersistentTrackID_Invalid)
        else { return nil }

        do {
            try videoCompositionTrack.insertTimeRange(
                CMTimeRange(start: .zero, duration: videoAsset.duration),
                of: videoTrack,
                at: .zero
            )
        } catch { return nil }

        // Get the audio track from the audio asset
        guard let audioTrack = audioAsset.tracks(withMediaType: .audio).first else { return nil }
        // Create an audio composition track
        guard let audioCompositionTrack = mixComposition.addMutableTrack(
            withMediaType: .audio,
            preferredTrackID: kCMPersistentTrackID_Invalid
        )
        else { return nil }

        do {
            let audioDuration = min(audioAsset.duration, videoAsset.duration) // Limit audio to video length
            try audioCompositionTrack.insertTimeRange(
                CMTimeRange(start: .zero, duration: audioDuration),
                of: audioTrack,
                at: .zero
            )
        } catch { return nil }

        // Export the composition
        guard
            let exportSession = AVAssetExportSession(
                asset: mixComposition,
                presetName: AVAssetExportPresetHEVC1920x1080
            )
        else { return nil }

        exportSession.outputURL = Adamantium.createRandomFileURLInDocuments("mp4")
        exportSession.outputFileType = .mp4
        await exportSession.export()

        switch exportSession.status {
        case .completed:
            return exportSession.outputURL
        default:
            return nil
        }
    }

    // Public method to merge video and audio files, limiting the audio to the video's length, and export to the specified output URL
    func mergeVideoAndAudio(
        videoURL: URL,
        audioURL: URL,
        outputURL: URL,
        completion: @escaping (Result<URL, Error>) -> Void
    ) async {
        // Create the composition
        let composition = AVMutableComposition()
        var videoATimeRange: CMTimeRange = .zero

        // Add video track
        do {
            try await addVideoTrack(
                to: composition,
                from: videoURL,
                videoTimeRange: &videoATimeRange
            )
        } catch {
            completion(.failure(error))
            return
        }

        // Add audio track, limiting the duration to the video duration
        do {
            try await addAudioTrack(
                to: composition,
                from: audioURL,
                withVideoTimeRange: videoATimeRange
            )
        } catch {
            completion(.failure(error))
            return
        }

        // Export the composition to the specified output URL
        exportComposition(composition, to: outputURL, completion: completion)
    }

    // Public method to concatenate two videos with no transition and export to the specified output URL
    func concatenateVideos(
        videoURLA: URL,
        videoURLB: URL,
        outputURL: URL,
        completion: @escaping (Result<URL, Error>) -> Void
    ) async {
        // Create the composition
        let composition = AVMutableComposition()
        var videoATimeRange: CMTimeRange = .zero
        var videoBTimeRange: CMTimeRange = .zero

        // Add first video track
        do {
            try await addVideoTrack(
                to: composition,
                from: videoURLA,
                videoTimeRange: &videoATimeRange
            )
        } catch {
            completion(.failure(error))
            return
        }

        // Add second video track, appending it to the end of the first video
        do {
            try await addVideoTrack(
                to: composition,
                from: videoURLB,
                at: composition.duration,
                videoTimeRange: &videoBTimeRange
            )
        } catch {
            completion(.failure(error))
            return
        }

        // Export the composition to the specified output URL
        exportComposition(composition, to: outputURL, completion: completion)
    }
}

// MARK: - Private Extension for Helper Methods

private extension VideoCompositionUtility {
    // Adds a video track to the composition from the provided video URL at a specified start time
    func addVideoTrack(
        to composition: AVMutableComposition,
        from videoURL: URL,
        at startTime: CMTime = .zero,
        videoTimeRange: inout CMTimeRange
    ) async throws {
        let asset = AVAsset(url: videoURL)
        guard
            let videoAsset = try await asset.loadTracks(withMediaType: .video).first
        else { throw VideoCompositionError.couldNotAddVideoTrack }

        let videoTrack = composition.addMutableTrack(
            withMediaType: .video,
            preferredTrackID: kCMPersistentTrackID_Invalid
        )

        // let videoDuration = videoAsset.timeRange.duration
        let timeRange = try await videoAsset.load(.timeRange)
        videoTimeRange = timeRange
        try videoTrack?.insertTimeRange(
            timeRange,
            of: videoAsset,
            at: startTime
        )
    }

    // Adds an audio track to the composition from the provided audio URL,
    // limiting the duration to match the video duration

    func addAudioTrack(
        to composition: AVMutableComposition,
        from audioURL: URL,
        withVideoTimeRange videoTimeRange: CMTimeRange
    ) async throws {
        let asset = AVAsset(url: audioURL)
        guard
            let audioAsset = try await asset.loadTracks(withMediaType: .audio).first
        else { throw VideoCompositionError.couldNotAddAudioTrack }

        let audioTrack = composition.addMutableTrack(
            withMediaType: .audio,
            preferredTrackID: kCMPersistentTrackID_Invalid
        )

        try audioTrack?.insertTimeRange(
            videoTimeRange,
            of: audioAsset,
            at: .zero
        )
    }

    // Exports the given composition to the specified output URL
    func exportComposition(
        _ composition: AVMutableComposition,
        to outputURL: URL,
        completion: @escaping (Result<URL, Error>) -> Void
    ) {
        guard
            let exportSession = AVAssetExportSession(
                asset: composition,
                presetName: AVAssetExportPresetHighestQuality
            )
        else {
            completion(.failure(VideoCompositionError.creationSessionFailed))
            return
        }

        exportSession.outputURL = outputURL
        exportSession.outputFileType = .mp4
        exportSession.shouldOptimizeForNetworkUse = true
        exportSession.exportAsynchronously {
            switch exportSession.status {
            case .completed:
                completion(.success(outputURL))

            case .failed:
                if let error = exportSession.error {
                    completion(.failure(error))
                }

            case .cancelled:
                completion(.failure(VideoCompositionError.compositionExportFailed))

            default:
                break
            }
        }
    }
}

// Usage example:
/*
    let videoURL1 = URL(fileURLWithPath: "/path/to/video1.mp4")
    let videoURL2 = URL(fileURLWithPath: "/path/to/video2.mp4")
    let outputURL = URL(fileURLWithPath: "/path/to/output.mp4")

    let videoComposition = VideoCompositionUtility()
    videoComposition.concatenateVideos(
        videoURLA: videoURL1,
        videoURLB: videoURL2,
        outputURL: outputURL) { result in
        switch result {
        case .success(let url):
            print("Successfully concatenated videos at: \(url)")
        case .failure(let error):
            print("Failed to concatenate videos: \(error)")
        }
    }
 */