#include "../Core/output.hpp"
#import <AppKit/AppKit.h>
#import <AudioToolbox/AudioToolbox.h>
#import <CoreAudio/CoreAudio.h>
#import <MetalKit/MetalKit.h>
#import <mach/mach_time.h>
#import <simd/simd.h>
extern "C" {
#include <libswscale/swscale.h>
}
#include <algorithm>
#include <array>
#include <atomic>
#include <cmath>
#include <cstring>
#include <mutex>
#include <stdexcept>
#include <vector>

#ifdef PLAYER_AUDIO_CHECK
// Check builds observe the actual callback samples; release builds have no hook.
extern void player_check_audio(const float *samples, size_t frames);
#endif

namespace {
constexpr size_t capacity = 48000;
double host_seconds(uint64_t ticks) {
    static const double factor = [] {
        mach_timebase_info_data_t info;
        mach_timebase_info(&info);
        return double(info.numer) / double(info.denom) / 1e9;
    }();
    return ticks * factor;
}
struct VideoMailbox {
    std::mutex mutex;
    AVFrame *frame = av_frame_alloc();
    uint64_t version = 0;
    ~VideoMailbox() { av_frame_free(&frame); }
};
struct Uniforms { simd_float4 red, green, blue, parameters; };
NSString *const shader = @R"metal(
#include <metal_stdlib>
using namespace metal;
struct U { float4 red, green, blue, parameters; };
struct V { float4 position [[position]]; float2 uv; };
vertex V vertex_main(uint index [[vertex_id]], constant U &u [[buffer(0)]]) {
    const float2 p[] = { {-1,-1}, {1,-1}, {-1,1}, {1,1} };
    V v;
    v.position = float4(p[index] * u.parameters.xy, 0, 1);
    v.uv = float2((p[index].x + 1) / 2, (1 - p[index].y) / 2);
    return v;
}
fragment float4 fragment_main(V v [[stage_in]], constant U &u [[buffer(0)]],
    texture2d<float> y [[texture(0)]], texture2d<float> cb [[texture(1)]],
    texture2d<float> cr [[texture(2)]]) {
    constexpr sampler s(address::clamp_to_edge, filter::linear);
    if (u.parameters.w == 2) return float4(y.sample(s, v.uv).rgb, 1);
    float luma = y.sample(s, v.uv).r * u.parameters.z;
    float2 chroma = u.parameters.w == 1 ? cb.sample(s, v.uv).rg :
        float2(cb.sample(s, v.uv).r, cr.sample(s, v.uv).r);
    float4 values = float4(luma, chroma * u.parameters.z, 1);
    return float4(dot(u.red, values), dot(u.green, values), dot(u.blue, values), 1);
}
)metal";
}

@interface PlaybackVideoView : MTKView <MTKViewDelegate> {
    std::shared_ptr<VideoMailbox> mailbox;
    uint64_t seen;
    id<MTLCommandQueue> commands;
    id<MTLRenderPipelineState> pipeline;
    id<MTLTexture> planes[3];
    Uniforms uniforms;
    int imageWidth, imageHeight;
    double imageAspect;
    SwsContext *conversion;
    AVFrame *converted;
    std::shared_ptr<std::atomic_bool> inFlight;
    CFStringRef colorName;
}
- (instancetype)initWithMailbox:(std::shared_ptr<VideoMailbox>)source;
@end

@implementation PlaybackVideoView
- (instancetype)initWithMailbox:(std::shared_ptr<VideoMailbox>)source {
    id<MTLDevice> device = MTLCreateSystemDefaultDevice();
    if (!device) throw std::runtime_error("Metal is unavailable on this Mac");
    self = [super initWithFrame:NSMakeRect(0, 0, 960, 540) device:device];
    if (!self) throw std::runtime_error("Could not create the video view");
    mailbox = std::move(source);
    seen = UINT64_MAX;
    inFlight = std::make_shared<std::atomic_bool>(false);
    self.colorPixelFormat = MTLPixelFormatBGRA8Unorm;
    self.clearColor = MTLClearColorMake(0, 0, 0, 1);
    self.preferredFramesPerSecond = 60;
    self.delegate = self;
    self.layer.contentsScale = NSScreen.mainScreen.backingScaleFactor;
    commands = [device newCommandQueue];
    NSError *error = nil;
    id<MTLLibrary> library = [device newLibraryWithSource:shader options:nil error:&error];
    if (!library) throw std::runtime_error(error.localizedDescription.UTF8String);
    MTLRenderPipelineDescriptor *description = [MTLRenderPipelineDescriptor new];
    description.vertexFunction = [library newFunctionWithName:@"vertex_main"];
    description.fragmentFunction = [library newFunctionWithName:@"fragment_main"];
    description.colorAttachments[0].pixelFormat = self.colorPixelFormat;
    pipeline = [device newRenderPipelineStateWithDescriptor:description error:&error];
    if (!pipeline || !commands) throw std::runtime_error("Could not initialize Metal rendering");
    converted = av_frame_alloc();
    if (!converted) throw std::bad_alloc();
    return self;
}
- (void)dealloc {
    self.delegate = nil;
    sws_freeContext(conversion);
    av_frame_free(&converted);
}
- (void)mtkView:(MTKView *)view drawableSizeWillChange:(CGSize)size {}
- (void)upload:(const AVFrame *)source {
    const AVFrame *frame = source;
    AVPixelFormat format = static_cast<AVPixelFormat>(frame->format);
    bool planar = format == AV_PIX_FMT_YUV420P || format == AV_PIX_FMT_YUVJ420P ||
                  format == AV_PIX_FMT_YUV420P10LE;
    bool nv12 = format == AV_PIX_FMT_NV12;
    bool rgba = format == AV_PIX_FMT_RGBA || format == AV_PIX_FMT_BGRA;
    bool negativeStride = frame->linesize[0] < 0 || (!rgba && frame->linesize[1] < 0);
    if ((!planar && !nv12 && !rgba) || negativeStride) {
        conversion = sws_getCachedContext(conversion, frame->width, frame->height, format,
            frame->width, frame->height, AV_PIX_FMT_BGRA, SWS_BILINEAR, nullptr, nullptr, nullptr);
        if (!conversion) throw std::runtime_error("Pixel conversion is unavailable");
        int space = frame->colorspace == AVCOL_SPC_BT709 ? SWS_CS_ITU709 :
                    frame->colorspace == AVCOL_SPC_BT2020_NCL ? SWS_CS_BT2020 : SWS_CS_ITU601;
        const int *coefficients = sws_getCoefficients(space);
        sws_setColorspaceDetails(conversion, coefficients, frame->color_range == AVCOL_RANGE_JPEG,
                                coefficients, 1, 0, 1 << 16, 1 << 16);
        if (!converted->buf[0] || converted->width != frame->width || converted->height != frame->height) {
            av_frame_unref(converted);
            converted->format = AV_PIX_FMT_BGRA;
            converted->width = frame->width;
            converted->height = frame->height;
            if (av_frame_get_buffer(converted, 32) < 0) throw std::bad_alloc();
        }
        if (av_frame_make_writable(converted) < 0) throw std::bad_alloc();
        sws_scale(conversion, frame->data, frame->linesize, 0, frame->height,
                  converted->data, converted->linesize);
        frame = converted;
        format = AV_PIX_FMT_BGRA;
        rgba = true; planar = nv12 = false;
    }
    bool ten = format == AV_PIX_FMT_YUV420P10LE;
    int count = rgba ? 1 : nv12 ? 2 : 3;
    for (int i = 0; i < count; ++i) {
        NSUInteger width = i == 0 ? frame->width : (frame->width + 1) / 2;
        NSUInteger height = i == 0 ? frame->height : (frame->height + 1) / 2;
        MTLPixelFormat pixel = rgba ? (format == AV_PIX_FMT_BGRA ? MTLPixelFormatBGRA8Unorm : MTLPixelFormatRGBA8Unorm) :
            ten ? MTLPixelFormatR16Unorm : nv12 && i == 1 ? MTLPixelFormatRG8Unorm : MTLPixelFormatR8Unorm;
        if (!planes[i] || planes[i].width != width || planes[i].height != height || planes[i].pixelFormat != pixel) {
            MTLTextureDescriptor *description = [MTLTextureDescriptor texture2DDescriptorWithPixelFormat:pixel
                width:width height:height mipmapped:NO];
            description.usage = MTLTextureUsageShaderRead;
            description.storageMode = MTLStorageModeShared;
            planes[i] = [self.device newTextureWithDescriptor:description];
            if (!planes[i]) throw std::runtime_error("Could not allocate video texture");
        }
        [planes[i] replaceRegion:MTLRegionMake2D(0, 0, width, height) mipmapLevel:0
                      withBytes:frame->data[i] bytesPerRow:frame->linesize[i]];
    }
    // Unused slots still need valid textures for the shared fragment function.
    if (count < 3) planes[2] = planes[count - 1];
    if (count == 1) planes[1] = planes[0];
    imageWidth = frame->width; imageHeight = frame->height;
    AVRational aspect = source->sample_aspect_ratio;
    imageAspect = double(imageWidth) / imageHeight * (aspect.num > 0 && aspect.den > 0 ? double(aspect.num) / aspect.den : 1);
    CFStringRef targetColor = source->color_primaries == AVCOL_PRI_BT2020 ? kCGColorSpaceITUR_2020 :
        source->color_trc == AVCOL_TRC_IEC61966_2_1 ? kCGColorSpaceSRGB : kCGColorSpaceITUR_709;
    if (colorName != targetColor) {
        CGColorSpaceRef space = CGColorSpaceCreateWithName(targetColor);
        self.colorspace = space;
        CGColorSpaceRelease(space);
        colorName = targetColor;
    }
    double kr = .299, kb = .114;
    if (source->colorspace == AVCOL_SPC_BT709 ||
        (source->colorspace == AVCOL_SPC_UNSPECIFIED && source->height >= 720)) { kr = .2126; kb = .0722; }
    if (source->colorspace == AVCOL_SPC_BT2020_NCL) { kr = .2627; kb = .0593; }
    bool full = source->color_range == AVCOL_RANGE_JPEG || format == AV_PIX_FMT_YUVJ420P;
    double ys = full ? 1 : ten ? 1023. / 876 : 255. / 219;
    double cs = full ? 1 : ten ? 1023. / 896 : 255. / 224;
    double offset = full ? 0 : ten ? 64. / 1023 : 16. / 255;
    double center = ten ? 512. / 1023 : 128. / 255;
    double r = 2 * (1 - kr) * cs, b = 2 * (1 - kb) * cs;
    double gc = -2 * kb * (1 - kb) / (1 - kr - kb) * cs;
    double gr = -2 * kr * (1 - kr) / (1 - kr - kb) * cs;
    uniforms.red = {float(ys), 0, float(r), float(-ys * offset - r * center)};
    uniforms.green = {float(ys), float(gc), float(gr), float(-ys * offset - (gc + gr) * center)};
    uniforms.blue = {float(ys), float(b), 0, float(-ys * offset - b * center)};
    uniforms.parameters = {1, 1, ten ? float(65535. / 1023) : 1.f, rgba ? 2.f : nv12 ? 1.f : 0.f};
}
- (void)drawInMTKView:(MTKView *)view {
    // Do not overwrite shared textures while the GPU is still sampling them.
    if (inFlight->exchange(true, std::memory_order_acquire)) return;
    // Keep the mailbox lock only for retaining a frame, never during GPU work.
    AVFrame *frame = nullptr;
    {
        std::lock_guard lock(mailbox->mutex);
        if (seen != mailbox->version) {
            seen = mailbox->version;
            if (mailbox->frame->buf[0]) frame = av_frame_clone(mailbox->frame);
            else { planes[0] = planes[1] = planes[2] = nil; imageWidth = imageHeight = 0; }
        }
    }
    if (frame) {
        try { [self upload:frame]; }
        catch (const std::exception &error) { NSLog(@"Video rendering failed: %s", error.what()); }
        av_frame_free(&frame);
    }
    MTLRenderPassDescriptor *pass = self.currentRenderPassDescriptor;
    id<CAMetalDrawable> drawable = self.currentDrawable;
    if (!pass || !drawable) { inFlight->store(false, std::memory_order_release); return; }
    id<MTLCommandBuffer> command = [commands commandBuffer];
    id<MTLRenderCommandEncoder> encoder = [command renderCommandEncoderWithDescriptor:pass];
    if (planes[0] && imageWidth > 0) {
        double viewAspect = self.drawableSize.width / std::max(1., self.drawableSize.height);
        uniforms.parameters.x = imageAspect < viewAspect ? imageAspect / viewAspect : 1;
        uniforms.parameters.y = imageAspect > viewAspect ? viewAspect / imageAspect : 1;
        [encoder setRenderPipelineState:pipeline];
        [encoder setVertexBytes:&uniforms length:sizeof(uniforms) atIndex:0];
        [encoder setFragmentBytes:&uniforms length:sizeof(uniforms) atIndex:0];
        for (int i = 0; i < 3; ++i) [encoder setFragmentTexture:planes[i] atIndex:i];
        [encoder drawPrimitives:MTLPrimitiveTypeTriangleStrip vertexStart:0 vertexCount:4];
    }
    [encoder endEncoding];
    auto flight = inFlight;
    [command addCompletedHandler:^(id<MTLCommandBuffer>) {
        flight->store(false, std::memory_order_release);
    }];
    [command presentDrawable:drawable];
    [command commit];
}
@end

namespace playback {
struct Output::Impl {
    AudioUnit unit = nullptr;
    std::array<float, capacity * 2> samples{};
    std::atomic<uint64_t> read{0}, written{0};
    double origin = 0;
    double latency = 0;
    std::atomic<uint64_t> clockVersion{0};
    std::atomic<double> anchorPTS{NAN}, anchorHost{0}, endPTS{0}, frozen{NAN};
    bool running = false;
    std::shared_ptr<VideoMailbox> video = std::make_shared<VideoMailbox>();

    static OSStatus render(void *context, AudioUnitRenderActionFlags *, const AudioTimeStamp *time,
                           UInt32, UInt32 count, AudioBufferList *buffers) {
        auto &self = *static_cast<Impl *>(context);
        for (UInt32 i = 0; i < buffers->mNumberBuffers; ++i)
            std::memset(buffers->mBuffers[i].mData, 0, buffers->mBuffers[i].mDataByteSize);
        uint64_t r = self.read.load(std::memory_order_relaxed);
        uint64_t w = self.written.load(std::memory_order_acquire);
        size_t frames = std::min<uint64_t>(count, w - r);
        float *destination = static_cast<float *>(buffers->mBuffers[0].mData);
        size_t first = std::min(frames, capacity - size_t(r % capacity));
        std::memcpy(destination, self.samples.data() + (r % capacity) * 2, first * 2 * sizeof(float));
        std::memcpy(destination + first * 2, self.samples.data(), (frames - first) * 2 * sizeof(float));
        #ifdef PLAYER_AUDIO_CHECK
        player_check_audio(destination, count);
        #endif
        if (frames) {
            self.clockVersion.fetch_add(1, std::memory_order_acq_rel);
            self.anchorPTS.store(self.origin + double(r) / 48000, std::memory_order_relaxed);
            uint64_t host = time->mFlags & kAudioTimeStampHostTimeValid ? time->mHostTime : mach_absolute_time();
            self.anchorHost.store(host_seconds(host) + self.latency, std::memory_order_relaxed);
            self.endPTS.store(self.origin + double(r + frames) / 48000, std::memory_order_relaxed);
            self.clockVersion.fetch_add(1, std::memory_order_release);
        }
        self.read.store(r + frames, std::memory_order_release);
        return noErr;
    }
    ~Impl() {
        if (unit) {
            AudioOutputUnitStop(unit);
            AudioUnitUninitialize(unit);
            AudioComponentInstanceDispose(unit);
        }
    }
};
Output::Output() : impl(std::make_unique<Impl>()) { host_seconds(mach_absolute_time()); }
Output::~Output() = default;
bool Output::configure_audio(std::string &error) {
    if (impl->unit) return true;
    AudioComponentDescription description{ kAudioUnitType_Output, kAudioUnitSubType_DefaultOutput,
                                          kAudioUnitManufacturer_Apple, 0, 0 };
    AudioComponent component = AudioComponentFindNext(nullptr, &description);
    OSStatus result = component ? AudioComponentInstanceNew(component, &impl->unit) : -1;
    AudioStreamBasicDescription format{};
    format.mSampleRate = 48000;
    format.mFormatID = kAudioFormatLinearPCM;
    format.mFormatFlags = kAudioFormatFlagsNativeFloatPacked;
    format.mBytesPerPacket = format.mBytesPerFrame = 2 * sizeof(float);
    format.mFramesPerPacket = 1; format.mChannelsPerFrame = 2; format.mBitsPerChannel = 32;
    AURenderCallbackStruct callback{Impl::render, impl.get()};
    if (result == noErr) result = AudioUnitSetProperty(impl->unit, kAudioUnitProperty_StreamFormat,
        kAudioUnitScope_Input, 0, &format, sizeof(format));
    if (result == noErr) result = AudioUnitSetProperty(impl->unit, kAudioUnitProperty_SetRenderCallback,
        kAudioUnitScope_Input, 0, &callback, sizeof(callback));
    if (result == noErr) result = AudioUnitInitialize(impl->unit);
    if (result != noErr) {
        error = "Core Audio initialization failed (" + std::to_string(result) + ")";
        if (impl->unit) AudioComponentInstanceDispose(impl->unit);
        impl->unit = nullptr;
        return false;
    }
    AudioDeviceID device = 0;
    UInt32 size = sizeof(device);
    if (AudioUnitGetProperty(impl->unit, kAudioOutputUnitProperty_CurrentDevice,
        kAudioUnitScope_Global, 0, &device, &size) == noErr) {
        AudioObjectPropertyAddress address{ kAudioDevicePropertyLatency, kAudioObjectPropertyScopeOutput,
                                            kAudioObjectPropertyElementMain };
        UInt32 latency = 0; size = sizeof(latency);
        Float64 rate = 48000;
        AudioObjectGetPropertyData(device, &address, 0, nullptr, &size, &latency);
        address.mSelector = kAudioDevicePropertyNominalSampleRate; address.mScope = kAudioObjectPropertyScopeGlobal;
        size = sizeof(rate);
        AudioObjectGetPropertyData(device, &address, 0, nullptr, &size, &rate);
        address = {kAudioDevicePropertyStreams, kAudioObjectPropertyScopeOutput, kAudioObjectPropertyElementMain};
        if (AudioObjectGetPropertyDataSize(device, &address, 0, nullptr, &size) == noErr) {
            std::vector<AudioStreamID> streams(size / sizeof(AudioStreamID));
            if (AudioObjectGetPropertyData(device, &address, 0, nullptr, &size, streams.data()) == noErr) {
                UInt32 streamLatency = 0;
                for (AudioStreamID stream : streams) {
                    AudioObjectPropertyAddress property{kAudioStreamPropertyLatency,
                        kAudioObjectPropertyScopeGlobal, kAudioObjectPropertyElementMain};
                    UInt32 value = 0, bytes = sizeof(value);
                    if (AudioObjectGetPropertyData(stream, &property, 0, nullptr, &bytes, &value) == noErr)
                        streamLatency = std::max(streamLatency, value);
                }
                latency += streamLatency;
            }
        }
        impl->latency = rate > 0 ? latency / rate : 0;
    }
    return true;
}
bool Output::start(std::string &error) {
    if (!impl->unit) return true; // Video-only playback does not open an audio device.
    if (impl->running) return true;
    OSStatus status = AudioOutputUnitStart(impl->unit);
    if (status != noErr) { error = "Could not start Core Audio (" + std::to_string(status) + ")"; return false; }
    impl->running = true;
    return true;
}
void Output::pause() {
    if (!impl->running) return;
    double position = clock();
    AudioOutputUnitStop(impl->unit);
    impl->running = false;
    impl->frozen.store(position, std::memory_order_release);
}
void Output::reset() {
    impl->read.store(0); impl->written.store(0);
    impl->origin = 0;
    impl->anchorPTS.store(NAN); impl->frozen.store(NAN); impl->endPTS.store(0);
    std::lock_guard lock(impl->video->mutex);
    av_frame_unref(impl->video->frame);
    ++impl->video->version;
}
size_t Output::write(const float *samples, size_t frames, double pts) {
    uint64_t w = impl->written.load(std::memory_order_relaxed);
    uint64_t r = impl->read.load(std::memory_order_acquire);
    frames = std::min(frames, capacity - size_t(w - r));
    if (!frames) return 0;
    if (w == 0) impl->origin = pts;
    size_t first = std::min(frames, capacity - size_t(w % capacity));
    std::memcpy(impl->samples.data() + (w % capacity) * 2, samples, first * 2 * sizeof(float));
    std::memcpy(impl->samples.data(), samples + first * 2, (frames - first) * 2 * sizeof(float));
    impl->written.store(w + frames, std::memory_order_release);
    return frames;
}
size_t Output::queued() const {
    return impl->written.load(std::memory_order_acquire) - impl->read.load(std::memory_order_acquire);
}
double Output::clock() const {
    if (!impl->running) return impl->frozen.load(std::memory_order_acquire);
    for (int attempt = 0; attempt < 3; ++attempt) {
        uint64_t version = impl->clockVersion.load(std::memory_order_acquire);
        if (version & 1) continue;
        double pts = impl->anchorPTS.load(std::memory_order_relaxed);
        double host = impl->anchorHost.load(std::memory_order_relaxed);
        double end = impl->endPTS.load(std::memory_order_relaxed);
        if (version != impl->clockVersion.load(std::memory_order_acquire)) continue;
        if (!std::isfinite(pts)) return NAN;
        return std::clamp(pts + host_seconds(mach_absolute_time()) - host, pts, end);
    }
    return NAN;
}
void Output::present(const AVFrame *frame) {
    std::lock_guard lock(impl->video->mutex);
    av_frame_unref(impl->video->frame);
    if (av_frame_ref(impl->video->frame, frame) < 0) throw std::bad_alloc();
    ++impl->video->version;
}
void *Output::make_video_view() {
    PlaybackVideoView *view = [[PlaybackVideoView alloc] initWithMailbox:impl->video];
    return (__bridge_retained void *)view;
}
}