#include "video_decoder.h"

#include <mfapi.h>
#include <mferror.h>
#include <mfidl.h>

using Microsoft::WRL::ComPtr;

// ─── Minimal SPS bit-reader ───────────────────────────────────────────────────

struct BitReader {
  const uint8_t* data;
  size_t         pos = 0;

  uint32_t Bit() {
    uint32_t v = (data[pos / 8] >> (7 - pos % 8)) & 1;
    ++pos;
    return v;
  }
  uint32_t Bits(int n) {
    uint32_t v = 0;
    for (int i = 0; i < n; ++i) v = (v << 1) | Bit();
    return v;
  }
  uint32_t UEG() {
    int lz = 0;
    while (Bit() == 0 && lz < 32) ++lz;
    if (lz == 0) return 0;
    return (1u << lz) - 1 + Bits(lz);
  }
  int32_t SEG() {
    uint32_t code = UEG();
    return (code & 1) ? static_cast<int32_t>((code + 1) / 2)
                      : -static_cast<int32_t>(code / 2);
  }
};

// ─── SPS parser ───────────────────────────────────────────────────────────────

bool VideoDecoder::ParseSpsSize(const uint8_t* data, size_t len,
                                int& out_w, int& out_h) {
  size_t nal_start = SIZE_MAX;
  for (size_t i = 0; i + 4 <= len; ++i) {
    if (data[i] == 0 && data[i+1] == 0 && data[i+2] == 0 && data[i+3] == 1 &&
        i + 4 < len && (data[i+4] & 0x1F) == 7) {
      nal_start = i + 4; break;
    }
    if (data[i] == 0 && data[i+1] == 0 && data[i+2] == 1 &&
        i + 3 < len && (data[i+3] & 0x1F) == 7) {
      nal_start = i + 3; break;
    }
  }
  if (nal_start == SIZE_MAX) return false;

  BitReader br{data + nal_start + 1};

  uint8_t profile_idc = static_cast<uint8_t>(br.Bits(8));
  br.Bits(8); br.Bits(8); br.UEG();

  static const uint8_t kHiProfiles[] = {100,110,122,244,44,83,86,118,128};
  bool hi = false;
  for (uint8_t p : kHiProfiles) if (profile_idc == p) { hi = true; break; }
  if (hi) {
    uint32_t chroma = br.UEG();
    if (chroma == 3) br.Bit();
    br.UEG(); br.UEG(); br.Bit();
    if (br.Bit()) {
      for (int idx = 0; idx < (chroma != 3 ? 8 : 12); ++idx) {
        if (br.Bit()) {
          int sz = idx < 6 ? 16 : 64;
          int last = 8, next = 8;
          for (int j = 0; j < sz; ++j) {
            if (next != 0) next = (last + br.SEG() + 256) % 256;
            last = next == 0 ? last : next;
          }
        }
      }
    }
  }

  br.UEG();
  uint32_t poc = br.UEG();
  if (poc == 0)      br.UEG();
  else if (poc == 1) { br.Bit(); br.SEG(); br.SEG(); uint32_t nc=br.UEG(); for(uint32_t j=0;j<nc;++j) br.SEG(); }

  br.UEG(); br.Bit();

  uint32_t pw = br.UEG();
  uint32_t ph = br.UEG();
  int w = (int)((pw + 1) * 16);
  int h = (int)((ph + 1) * 16);

  bool frame_mbs_only = br.Bit() != 0;
  if (!frame_mbs_only) br.Bit();
  br.Bit();

  if (br.Bit()) {
    uint32_t cl=br.UEG(), cr=br.UEG(), ct=br.UEG(), cb=br.UEG();
    int ux = 2, uy = frame_mbs_only ? 2 : 4;
    w -= ux * (int)(cl + cr);
    h -= uy * (int)(ct + cb);
  }

  if (w <= 0 || h <= 0) return false;
  out_w = w; out_h = h;
  return true;
}

// ─── Pixel buffer callback (called from render thread) ────────────────────────

const FlutterDesktopPixelBuffer* VideoDecoder::PixelBufferCallback(
    size_t, size_t, void* user_data) {
  return &static_cast<VideoDecoder*>(user_data)->pixel_buf_;
}

// ─── VideoDecoder ─────────────────────────────────────────────────────────────

VideoDecoder::VideoDecoder(FlutterDesktopTextureRegistrarRef tex_reg,
                           const uint8_t* sps_pps, size_t sps_pps_len)
    : tex_reg_(tex_reg) {
  Init(sps_pps, sps_pps_len);
}

VideoDecoder::~VideoDecoder() {
  if (texture_id_ >= 0) {
    HANDLE done = CreateEvent(nullptr, TRUE, FALSE, nullptr);
    FlutterDesktopTextureRegistrarUnregisterExternalTexture(
        tex_reg_, texture_id_,
        [](void* ud) { SetEvent(static_cast<HANDLE>(ud)); },
        done);
    WaitForSingleObject(done, 3000);
    CloseHandle(done);
  }
  if (decoder_)
    decoder_->ProcessMessage(MFT_MESSAGE_NOTIFY_END_OF_STREAM, 0);
}

bool VideoDecoder::Init(const uint8_t* sps_pps, size_t sps_pps_len) {
  int w = 0, h = 0;
  if (!ParseSpsSize(sps_pps, sps_pps_len, w, h)) {
    w = GetSystemMetrics(SM_CXSCREEN);
    h = GetSystemMetrics(SM_CYSCREEN);
  }
  width_ = w; height_ = h;

  MFT_REGISTER_TYPE_INFO in_info{MFMediaType_Video, MFVideoFormat_H264};
  UINT32 count = 0; IMFActivate** acts = nullptr;
  HRESULT hr = MFTEnumEx(MFT_CATEGORY_VIDEO_DECODER,
      MFT_ENUM_FLAG_SYNCMFT | MFT_ENUM_FLAG_SORTANDFILTER,
      &in_info, nullptr, &acts, &count);
  if (FAILED(hr) || count == 0) return false;
  hr = acts[0]->ActivateObject(IID_PPV_ARGS(&decoder_));
  for (UINT32 i = 0; i < count; ++i) acts[i]->Release();
  CoTaskMemFree(acts);
  if (FAILED(hr)) return false;

  ComPtr<IMFMediaType> in_type;
  MFCreateMediaType(&in_type);
  in_type->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Video);
  in_type->SetGUID(MF_MT_SUBTYPE,    MFVideoFormat_H264);
  MFSetAttributeSize(in_type.Get(), MF_MT_FRAME_SIZE, (UINT32)w, (UINT32)h);
  if (FAILED(decoder_->SetInputType(0, in_type.Get(), 0))) return false;

  ComPtr<IMFMediaType> out_type;
  MFCreateMediaType(&out_type);
  out_type->SetGUID(MF_MT_MAJOR_TYPE, MFMediaType_Video);
  out_type->SetGUID(MF_MT_SUBTYPE,    MFVideoFormat_NV12);
  MFSetAttributeSize(out_type.Get(), MF_MT_FRAME_SIZE, (UINT32)w, (UINT32)h);
  if (FAILED(decoder_->SetOutputType(0, out_type.Get(), 0))) return false;

  decoder_->ProcessMessage(MFT_MESSAGE_NOTIFY_BEGIN_STREAMING, 0);
  decoder_->ProcessMessage(MFT_MESSAGE_NOTIFY_START_OF_STREAM, 0);

  // Register Flutter BGRA texture via C API
  frame_data_.assign(static_cast<size_t>(w * h * 4), 0);
  pixel_buf_.buffer  = frame_data_.data();
  pixel_buf_.width   = static_cast<size_t>(w);
  pixel_buf_.height  = static_cast<size_t>(h);

  FlutterDesktopTextureInfo info{};
  info.type = kFlutterDesktopPixelBufferTexture;
  info.pixel_buffer_config.callback  = &VideoDecoder::PixelBufferCallback;
  info.pixel_buffer_config.user_data = this;

  texture_id_ = FlutterDesktopTextureRegistrarRegisterExternalTexture(
      tex_reg_, &info);

  // Prime decoder with SPS+PPS
  FeedNal(sps_pps, sps_pps_len);
  return texture_id_ >= 0;
}

// ─── NV12 → BGRA ──────────────────────────────────────────────────────────────

void VideoDecoder::Nv12ToBgra(const uint8_t* nv12, int stride,
                               uint8_t* bgra, int w, int h) {
  const uint8_t* Y  = nv12;
  const uint8_t* UV = nv12 + stride * h;

  for (int row = 0; row < h; ++row) {
    for (int col = 0; col < w; ++col) {
      int y  = Y[row * stride + col];
      int uv = (row / 2) * stride + (col & ~1);
      int u  = UV[uv];
      int v  = UV[uv + 1];

      int c = y - 16, d = u - 128, e = v - 128;
      int r = (298*c + 409*e + 128) >> 8;
      int g = (298*c - 100*d - 208*e + 128) >> 8;
      int b = (298*c + 516*d + 128) >> 8;

      uint8_t* p = bgra + (row * w + col) * 4;
      p[0] = static_cast<uint8_t>(b < 0 ? 0 : b > 255 ? 255 : b);
      p[1] = static_cast<uint8_t>(g < 0 ? 0 : g > 255 ? 255 : g);
      p[2] = static_cast<uint8_t>(r < 0 ? 0 : r > 255 ? 255 : r);
      p[3] = 255;
    }
  }
}

// ─── FeedNal ──────────────────────────────────────────────────────────────────

void VideoDecoder::FeedNal(const uint8_t* nal, size_t len) {
  std::lock_guard<std::mutex> lk(mu_);
  if (!decoder_ || !len) return;

  ComPtr<IMFSample>      in_sample;
  ComPtr<IMFMediaBuffer> in_buf;
  MFCreateMemoryBuffer((DWORD)len, &in_buf);
  {
    BYTE* ptr = nullptr;
    in_buf->Lock(&ptr, nullptr, nullptr);
    memcpy(ptr, nal, len);
    in_buf->Unlock();
  }
  in_buf->SetCurrentLength((DWORD)len);
  MFCreateSample(&in_sample);
  in_sample->SetSampleTime(input_ts_);
  in_sample->SetSampleDuration(333333);
  in_sample->AddBuffer(in_buf.Get());
  input_ts_ += 333333;

  if (FAILED(decoder_->ProcessInput(0, in_sample.Get(), 0))) return;

  MFT_OUTPUT_STREAM_INFO si{};
  decoder_->GetOutputStreamInfo(0, &si);

  while (true) {
    ComPtr<IMFSample>      out_sample;
    ComPtr<IMFMediaBuffer> out_buf;
    if (!(si.dwFlags & MFT_OUTPUT_STREAM_PROVIDES_SAMPLES)) {
      DWORD sz = si.cbSize ? si.cbSize : static_cast<DWORD>(width_ * height_ * 3 / 2);
      MFCreateMemoryBuffer(sz, &out_buf);
      MFCreateSample(&out_sample);
      out_sample->AddBuffer(out_buf.Get());
    }

    MFT_OUTPUT_DATA_BUFFER out_data{};
    out_data.pSample = out_sample.Get();
    DWORD   status = 0;
    HRESULT hr = decoder_->ProcessOutput(0, 1, &out_data, &status);
    if (out_data.pEvents) out_data.pEvents->Release();
    if (hr == MF_E_TRANSFORM_NEED_MORE_INPUT) break;
    if (FAILED(hr)) break;

    ComPtr<IMFSample> result(out_data.pSample);
    if (!result) break;

    ComPtr<IMFMediaBuffer> flat;
    result->ConvertToContiguousBuffer(&flat);
    BYTE*  nv12_ptr = nullptr; DWORD nv12_len = 0;
    flat->Lock(&nv12_ptr, nullptr, &nv12_len);
    Nv12ToBgra(nv12_ptr, width_, frame_data_.data(), width_, height_);
    flat->Unlock();

    FlutterDesktopTextureRegistrarMarkExternalTextureFrameAvailable(
        tex_reg_, texture_id_);
  }
}