import AVFoundation
import Foundation
import UIKit

func createLoopingVideoWithAudio(
    audioURL: URL,
    videoURL: URL,
    videoSize: CGSize,
    duration: Double // seconds
) async throws -> URL {
    let cachesDirectory = try FileManager.default.url(
        for: .cachesDirectory,
        in: .userDomainMask,
        appropriateFor: nil,
        create: false
    )

    let outputURL = cachesDirectory.appendingPathComponent("shareVideo").appendingPathExtension("mp4")
    try? FileManager.default.removeItem(at: outputURL)

    let audioAsset = AVAsset(url: audioURL)
    let audioTracks = try await audioAsset.loadTracks(withMediaType: .audio)
    let videoAsset = AVAsset(url: videoURL)
    let videoTracks = try await videoAsset.loadTracks(withMediaType: .video)
    let originalDuration = try await videoAsset.load(.duration)
    let requiredDuration = CMTime(seconds: duration, preferredTimescale: originalDuration.timescale)

    guard let videoTrack = videoTracks.first else {
        throw CreateMediaError.videoFileEmpty
    }

    let composition = AVMutableComposition()
    composition.naturalSize = .init(width: videoSize.width, height: videoSize.height)
    let videoCompositionTrack = composition.addMutableTrack(withMediaType: .video, preferredTrackID: kCMPersistentTrackID_Invalid)
    videoCompositionTrack?.preferredTransform = try await videoTrack.load(.preferredTransform)

    do {
        var currentTime = CMTime.zero

        while currentTime < requiredDuration {
            let timeRange = CMTimeRange(start: CMTime.zero, duration: originalDuration)

            try videoCompositionTrack?.insertTimeRange(timeRange, of: videoTrack, at: currentTime)

            currentTime = CMTimeAdd(currentTime, originalDuration)
        }

        for audioTrack in audioTracks {
            let audioCompositionTrack = composition.addMutableTrack(withMediaType: .audio, preferredTrackID: kCMPersistentTrackID_Invalid)
            try audioCompositionTrack?
                .insertTimeRange(CMTimeRange(start: CMTime.zero, end: requiredDuration), of: audioTrack, at: CMTime.zero)
        }

        let exportSession = AVAssetExportSession(asset: composition, presetName: AVAssetExportPresetHighestQuality)
        exportSession?.outputFileType = .mp4
        exportSession?.shouldOptimizeForNetworkUse = true
        exportSession?.outputURL = outputURL
        await exportSession?.export()

        return outputURL
    } catch {
        throw CreateMediaError.failedAVMediaComposition
    }
}