import SwiftUI
import Metal
import MetalKit
import CoreImage
import AVKit

protocol PreviewMetalViewCoordinatorDelegate: AnyObject {
    func clearAllOutputURLs()
    func setVideoOutput(_ url: URL?)
    func setAudioOutput(_ url: URL?)
    func setImageOutput(_ url: URL?)
    func setFrameRateUpdate(_ fps: Float)
    func didUpdateVideoRecordingFrame(_ deltaTime: TimeInterval, recordingState: CaptureRecordingState)
}

class PreviewMetalViewCoordinator: ObservableObject {
    
    weak var previewMetalView: PreviewMetalView?
    weak var audioCaptureCoordinator: AudioCaptureCoordinator?
    weak var videoCaptureCoordinator: VideoCaptureCoordinator?
    weak var delegate: PreviewMetalViewCoordinatorDelegate?
    
    private let recordingStateQueue = DispatchQueue(
        label: "RecordingStateQueue",
        attributes: [],
        autoreleaseFrequency: .workItem)
    
    @Published private(set) var recordingState: CaptureRecordingState = .needsSetup
    private var internalRecordingState: CaptureRecordingState = .needsSetup
    
    private var videoCaptureUtility: VideoCaptureUtility?
    private var videoCaptureFrameIndex: Int = .zero
    private var videoCaptureStartTime: Date = .now
    
    private var gpu: MTLDevice? {
        return Adamantium.sharedDevice
    }
    
    private var commandQueue: MTLCommandQueue?
    
    private let startTime: Date = .now
    private var lastTime: Date?

    private var passthroughPass = PassthroughRenderPass()
    private var segmentationRenderTarget: MTLTexture?
    
    init() {
        configureMetal()
    }
    
    required init(coder: NSCoder) {
        fatalError("init(coder:) has not been implemented")
    }
    
    func resetAfterActiveSession() {
        
        // Results
        delegate?.clearAllOutputURLs()
        
        // Global
        setRecordingState(.needsSetup)
        
        // Video
        videoCaptureFrameIndex = .zero
        setupVideoCaptureIfNeeded()
    }
    
    func configureMetal() {
        guard let gpu else { return }
        passthroughPass.setupRenderPass()
        commandQueue = gpu.makeCommandQueue()
    }
    
    func setPassthroughPixelBuffer(_ buffer: CVPixelBuffer?) {
        passthroughPass.pixelBuffer = buffer
    }
    
    func setPassthroughMirroring(_ isMirroring: Bool) {
        passthroughPass.mirroring = isMirroring
    }
    
    func setPassthroughRotation(_ rotation: PassthroughRenderPass.Rotation) {
        passthroughPass.rotation = rotation
    }
    
    func setPassthroughScale(_ scale: CGFloat) {
        passthroughPass.scale = scale
    }
    
    func capturePhoto()  {
        let recordingState = getRecordingState()
        guard recordingState == .readyToRecord else { return }
        setRecordingState(.isFlaggedForPhotoCaptureOnNextDraw)
    }
    
    func startRecording() {
        let recordingState = getRecordingState()
        guard recordingState == .readyToRecord else { return }
        setRecordingState(.isRecording(.zero, frameIndex: .zero))
    }
    
    func endRecording() {
        let recordingState = getRecordingState()
        // number doesn't matter for equal on this enum specifically
        guard recordingState == .isRecording(.zero, frameIndex: .zero) else { return }
        setRecordingState(.isFlaggedToCompleteRecordingOnNextDraw)
    }
    
    /// Called once per frame
    func draw(
        _ rect: CGRect,
        bounds: CGRect,
        drawable: CAMetalDrawable?,
        renderPassDescriptor: MTLRenderPassDescriptor?) {
            
        let recordingState = getRecordingState()
        if recordingState == .needsSetup, let drawable, let _ = gpu {
            setupVideoCaptureIfNeeded()
            warmUpBufferAllocator(drawable)
            setRecordingState(.readyToRecord)
        }
            
        let deltaTime = updateDeltaTime()

        guard
            let drawable,
            let renderPassDescriptor,
            let commandQueue = Adamantium.sharedCommandQueue,
            let commandBuffer = commandQueue.makeCommandBuffer(),
            let renderCommandEncoder = commandBuffer.makeRenderCommandEncoder(descriptor: renderPassDescriptor)
        else {
            print("Failed to create Metal renderCommandEncoder")
            passthroughPass.flushTextureCache()
            return
        }
        
        renderCommandEncoder.label = "Passthrough"
        passthroughPass.render(encoder: renderCommandEncoder, viewBounds: bounds)
        renderCommandEncoder.endEncoding()
        
        let texture = drawable.texture
        commandBuffer.addCompletedHandler { [weak self] _ in
            self?.updateRecordingFromState(texture, state: recordingState, deltaTime: deltaTime)
        }
        
        // Draw to the screen.
        commandBuffer.present(drawable)
        commandBuffer.commit()
    }
}

extension PreviewMetalViewCoordinator: PassthroughTextureViewable {
    func texturePointForView(point: CGPoint) -> CGPoint? {
        return passthroughPass.texturePointForView(point: point)
    }
    
    func viewPointForTexture(point: CGPoint) -> CGPoint? {
        return passthroughPass.viewPointForTexture(point: point)
    }
}

private extension PreviewMetalViewCoordinator {
    
    func setRecordingState(_ state: CaptureRecordingState) {
        DispatchQueue.main.async {
            self.recordingState = state
        }
        
        recordingStateQueue.sync {
            internalRecordingState = state
        }
    }
    
    func getRecordingState() -> CaptureRecordingState {
        var localRecordingState: CaptureRecordingState = .needsSetup
        
        recordingStateQueue.sync {
            localRecordingState = internalRecordingState
        }
        
        return localRecordingState
    }
    
    func updateDeltaTime() -> TimeInterval {
        let deltaTime: TimeInterval
        if let lastTime = lastTime {
            deltaTime = Date.now.timeIntervalSince(lastTime)
        } else {
            deltaTime = .zero
        }
        
        lastTime = .now
        
        return deltaTime
    }
    
    func updateRecordingFromState(
        _ texture: MTLTexture,
        state: CaptureRecordingState,
        deltaTime: TimeInterval) {
        
        var stateToSendToDelegate: CaptureRecordingState = state
            
        switch state {
        case .needsSetup, 
            .readyToRecord,
            .isEndingRecording,
            .completed:
            break
            
        case .isFlaggedForPhotoCaptureOnNextDraw:
            do {
                let url = try StillCaptureUtility.png(texture: texture)
                DispatchQueue.main.async { [weak self] in
                    self?.delegate?.setImageOutput(url)
                    self?.setRecordingState(.completed)
                }
            } catch {
                print(error)
            }
        case .isRecording:
            audioCaptureCoordinator?.startRecordingIfNeeded()
            addFrameToRecording(texture, sourceDeltaTime: deltaTime)
            let videoTimeElapsed = Date.now.timeIntervalSince(videoCaptureStartTime)
            stateToSendToDelegate = .isRecording(videoTimeElapsed, frameIndex: videoCaptureFrameIndex)
            
        case .isFlaggedToCompleteRecordingOnNextDraw:
            audioCaptureCoordinator?.stopRecording()
            endRecordingOnCaptureUtility()
        }
        
        DispatchQueue.main.async { [weak self] in
            self?.delegate?.didUpdateVideoRecordingFrame(deltaTime, recordingState: stateToSendToDelegate)
        }
    }
    
    func addFrameToRecording(_ texture: MTLTexture, sourceDeltaTime: TimeInterval) {
        guard let videoCaptureUtility else { return }
        if videoCaptureFrameIndex == .zero {
            videoCaptureStartTime = .now
        }
        /*
            The reason this looks so funky is if I use the calculated framerate
            from the rounding the sourceDelta corrupts the video when I use it in a time marker.
         
            However if I use the number literal version as seen below
            (i.e videoFrameDeltaTime = 1.0 / 60.0)
            it works as expected.
         
        */
        
        let frameRate = sourceDeltaTime > 0.0 ? (1.0 / sourceDeltaTime) : 60.0
        let adjustedFrameRate = (frameRate / 10.0).rounded() * 10.0
        delegate?.setFrameRateUpdate(Float(adjustedFrameRate))
        let videoFrameDeltaTime: TimeInterval = 1.0 / 60.0
        let customTimeElapsed = TimeInterval(videoCaptureFrameIndex) * videoFrameDeltaTime
        
        videoCaptureUtility.writeFrame(
            forTexture: texture,
            customTimeIncrease: customTimeElapsed,
            isBufferAllocationOnly: false)
        
        let nextFrameIndex = videoCaptureFrameIndex + 1
        if nextFrameIndex < frameLimit {
            videoCaptureFrameIndex = nextFrameIndex
        } else {
            setRecordingState(.isFlaggedToCompleteRecordingOnNextDraw)
        }
    }
    
    func endRecordingOnCaptureUtility() {
        guard let videoCaptureUtility else { return }
        setRecordingState(.isEndingRecording)
        videoCaptureUtility
            .endRecording { outputURL in
                DispatchQueue.main.async {
                    self.setRecordingState(.completed)
                    self.delegate?.setVideoOutput(outputURL)
                }
            }
    }
    
    func setupVideoCaptureIfNeeded() {
        guard
            let previewMetalView,
            let outputURL = Adamantium.createRandomFileURLInDocuments("mp4"),
            let newVideoCaptureUtility = VideoCaptureUtility(
                outputURL: outputURL,
                size: previewMetalView.drawableSize)
        else { return }
        
        newVideoCaptureUtility.startRecording()
        
        videoCaptureFrameIndex = .zero
        videoCaptureUtility = newVideoCaptureUtility
    }
    
    func warmUpBufferAllocator(_ drawable: CAMetalDrawable?) {
        guard let texture = drawable?.texture else { return }
        // If we don't do this there is a delay when we
        // start adding frames to the capture context
        videoCaptureUtility?.writeFrame(forTexture: texture, isBufferAllocationOnly: true)
    }
    
    var frameLimit: Int {
        // 30s Time Limit
        // Frame Rate (FPS) * Time (Seconds)
        return 60 * 30
    }
}
