import Foundation
import AVFoundation
import CoreImage
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)")
        }
    }
 */