import AVFoundation
import Foundation

public extension AVAsset {
    // Creates and exports a trimmed video in the input video's original quality.
    // This is used to show the user their selected input video in the same quality
    // it was imported in while trimming (and eventually editing).
    func trimmedForPreview(
        to outputUrl: URL,
        startTime: CMTime,
        endTime: CMTime
    ) async throws -> URL {
        try await trimmed(to: outputUrl, startTime: startTime, endTime: endTime, preset: AVAssetExportPreset1920x1080)
    }

    // Creates and exports a trimmed, compressed and upload-friendly video.
    func trimmedForExport(
        to outputUrl: URL,
        startTime: CMTime,
        endTime: CMTime
    ) async throws -> URL {
        #if targetEnvironment(simulator)
            try await trimmed(to: outputUrl, startTime: startTime, endTime: endTime, preset: AVAssetExportPreset1920x1080)
        #else
            try await trimmed(to: outputUrl, startTime: startTime, endTime: endTime, preset: AVAssetExportPreset1920x1080)
        #endif
    }

    private func trimmed(
        to outputUrl: URL,
        startTime: CMTime,
        endTime: CMTime,
        preset: String,
        keepOriginalAudio: Bool = false
    ) async throws -> URL {
        let duration = CMTimeSubtract(endTime, startTime)
        let timeRange = CMTimeRange(start: startTime, duration: duration)

        let composition = AVMutableComposition()

        do {
            let tracks = try await loadTracks(withMediaType: .video)
            guard let videoTrack = tracks.first else {
                throw TrimError.videoTrackMissing
            }

            let compositionTrack = composition.addMutableTrack(withMediaType: .video, preferredTrackID: kCMPersistentTrackID_Invalid)
            try compositionTrack?.insertTimeRange(timeRange, of: videoTrack, at: .zero)

            compositionTrack?.preferredTransform = try await videoTrack.load(.preferredTransform)

            if let audioTrack = try await loadTracks(withMediaType: .audio).first, keepOriginalAudio {
                let audioCompositionTrack = composition.addMutableTrack(withMediaType: .audio, preferredTrackID: kCMPersistentTrackID_Invalid)
                try audioCompositionTrack?.insertTimeRange(timeRange, of: audioTrack, at: .zero)
            }

            composition.naturalSize = try await videoTrack.load(.naturalSize)

            let exportSession = AVAssetExportSession(asset: composition, presetName: preset)
            exportSession?.outputURL = outputUrl
            exportSession?.outputFileType = .mp4
            exportSession?.timeRange = CMTimeRange(start: .zero, duration: duration)
            exportSession?.shouldOptimizeForNetworkUse = true

            await exportSession?.export()

            guard exportSession?.status == .completed else {
                throw TrimError.exportSessionError(exportSession?.error?.localizedDescription ?? "Unknown error")
            }

            return outputUrl
        } catch {
            throw error
        }
    }
}

public enum TrimError: LocalizedError {
    case videoTrackMissing
    case exportSessionError(String)

    public var errorDescription: String? {
        switch self {
        case .videoTrackMissing: "Video file corrupted"
        case .exportSessionError: "Invalid share url"
        }
    }

    var error: String? {
        switch self {
        case .videoTrackMissing: "Video track missing"
        case .exportSessionError(let message): "Export session failed with error: \(message)"
        }
    }
}