#include <algorithm>
#include <cstring>

#include "common.hh"
#include "util.hh"
#include "IffFile.hh"
#include "decompress.hh"
#include "MveDecoder.hh"

MveDecoder::MveDecoder(uint8_t const* base, size_t length)
  : iff_(base, length), width_(320), height_(165)
{
  //iff_.printStructure();

  auto& root = iff_.getRoot();
  if (!root.isForm())
    throw FormatException{"Root node not FORM"};

  auto& rootForm = dynamic_cast<IffFile::Form const&>(root);
  if (rootForm.getSubtype() != "MOVE")
    throw FormatException{"Root form not MOVE"};

  auto it = rootForm.childrenBegin();

  // Parse _PC_ chunk
  auto& PC = *it++;
  if ((PC.getType() != "_PC_") ||
      (PC.getSize() != 12))
    throw FormatException{"_PC_ chunk missing or wrong size"};

  const uint32_t PC_PALTcount = readU32LE(PC.begin()+8);
#ifndef NDEBUG
  printf("_PC_: %u PALTs\n", PC_PALTcount);
#endif

  // // Parse SOND chunk
  // auto& SOND = *it++;
  // if ((SOND.getType() != "SOND") ||
  //     (SOND.getSize() != 4))
  //   throw FormatException{"SOND chunk missing or wrong size"};

  // const uint32_t  SONDc = readU32LE(SOND.begin());
  // if (SONDc != 3)
  //   throw FormatException{"Unexpected SOND value"};

  // Parse data
  unsigned curPALT = 0, curBRCH = 0;
  bool haveIndex = false;
  // Store branch offsets, resolve to shots later
  std::vector<uint32_t> branches;
  Shot *curSHOT = nullptr;
  
  for (;it != rootForm.childrenEnd();++it) {
    if (it->getType() == "SIZE") {
      if (it->getSize() != 8)
	throw FormatException{"Unexpected SIZE size"};

      width_ = readU32LE(it->begin());
      height_ = readU32LE(it->begin()+4); 
    } else if (it->getType() == "PALT") {
      if (curPALT > PC_PALTcount)
	throw FormatException{"Number of PALT chunks exceeds amount specified in _PC_"};
      if (it->getSize() != 768)
	throw FormatException{"Unexpected PALT size"};
      palts_.emplace_back();
      
      std::transform(it->begin(), it->end(), palts_.back().begin(),
		     [](const char& in) -> uint8_t {return (in << 2) | ((in >> 6)&0x3);});
      ++curPALT;
    } else if (it->getType() == "SHOT") {
      if (!haveIndex) {
	branches_.resize(1);
	curBRCH = 0;
      }
      
      const uint32_t SHOTc = readU32LE(it->begin());
      if (SHOTc > PC_PALTcount)
	throw FormatException{"SHOT refers to PALT outside amount specified in _PC_"};
      Shot shot{palts_.at(SHOTc), {}, {}};
      branches_[curBRCH].push_back(shot);
      curSHOT = &branches_[curBRCH].back();
    } else if (it->getType() == "VGA ") {
      if (!curSHOT)
	throw FormatException{"VGA outside SHOT"};
      curSHOT->VGAs.push_back(&*it);
    } else if (it->getType() == "AUDI") {
      if (!curSHOT)
	throw FormatException{"AUDI outside SHOT"};
      curSHOT->AUDIs.push_back(&*it);
    } else if (it->getType() == "INDX") {
      if (haveIndex)
	throw FormatException{"Multiple INDX"};
      if (curSHOT)
	throw FormatException{"INDX after SHOT"};
      haveIndex = true;

      for(unsigned i = 0;i < it->getSize()/4;++i)
	branches.push_back(readU32LE(it->begin()+i*4));
      branches_.resize(branches.size());
    } else if (it->getType() == "BRCH") {
      size_t ofs = it->begin()-base-8;
      auto idxIt = std::find(branches.begin(), branches.end(), ofs);
      if (idxIt == branches.end())
	throw FormatException{"Could not resolve branch " + std::to_string(ofs)};
      curBRCH = (idxIt-branches.begin());
    } else if (it->getType() == "TEXT") {
      // Subtitle NYI
    } else if (it->getType() == "SOND") {
      // ignore
    } else
      throw FormatException{"Encountered unexpected chunk: " + it->getType()};
  }

#ifdef VIDDEBUG_ 
  unsigned i = 0;
#endif
  for (auto& brch : branches_) {
    for (auto& shot : brch) {
#ifdef VIDDEBUG_
      printf("Shot %u: Palette %ld, %lu video frames, %lu audio blocks\n",
	     i++, std::find(palts_.begin(), palts_.end(), shot.palt)-palts_.begin(),
	     shot.VGAs.size(), shot.AUDIs.size());
#endif
      if ((shot.AUDIs.size() != 0) &&
	  (shot.VGAs.size() != shot.AUDIs.size()))
      	throw FormatException{"Video/audio block count mismatch: V " +
	    std::to_string(shot.VGAs.size()) + ", A " +
	    std::to_string(shot.AUDIs.size())};
    }
  }
}

std::vector<char> parseHuff_(uint8_t const* data, size_t len)
{
  uint8_t numHuffVals;
  memcpy(&numHuffVals, data, 1);
  if (numHuffVals != 22)
    throw FormatException{"Unexpected huffman tree size"};
  
  std::array<uint8_t, 44> huffTree;
  memcpy(huffTree.data(), data+1, 44);

#ifdef HUFFDEBUG_
  printf("Coded data length: %d\nTree:\n", segOfs[1]-(segOfs[0]+45));
  printf("      00 01 02 03 04 05 06 07 08 09 0a 0b 0c 0d 0e 0f\n");
  printf("0x00: ");
  for(unsigned i = 0;i < huffTree.size();++i) {
    printf("%.2hhx", huffTree[i]);
    if ((i+1)%16 == 0)
      printf("\n0x%.2hhx: ", (i+1)&0xf0);
    else
      printf(" ");
  }
  printf("\n");
#endif

  int byteValid = 0;
  uint8_t byteBuf = 0;
  unsigned bytePos = 45;
  std::vector<char> commands;
  unsigned huffIdx = huffTree.size();
  while (huffIdx != 22) {
    // Read next bit
    unsigned bit;
    if (byteValid) {
      bit = byteBuf&0x1;
      byteBuf >>= 1;
      --byteValid;
    } else {
      if (bytePos >= len)
	throw FormatException{"Huffman stream overrun"};
      byteBuf = data[bytePos++];
      bit = byteBuf&0x1;
      byteBuf >>= 1;
      byteValid = 7;
    }

    huffIdx = huffTree.at(huffIdx-(bit?1:23));


    if (huffIdx < 22) {
      commands.push_back(huffIdx);
      huffIdx = huffTree.size();
    }
  }

  if (commands.empty())
    throw FormatException{"No commands decoded"};
  
#ifdef HUFFDEBUG_
  for(auto ent : commands)
    printf("%.2hhx ", ent);
  printf("\n");
#endif

  return commands;
}

std::vector<uint8_t> parsePixels_(uint8_t const* data, size_t len)
{
  if (*data == 0x02) {
    return decompressLZ(data+1, len-1);
  } else {
    std::vector<uint8_t> ret;
    std::copy(data+1, data+len, std::back_inserter(ret));
    return ret;
  }
}

MveDecoder::Movie MveDecoder::open(size_t branch) const
{
  auto& br = branches_.at(branch);

  return Movie(br.cbegin(), br.cend(),
	       *this);
}

MveDecoder::Movie::Movie(std::vector<Shot>::const_iterator shotIt,
			 std::vector<Shot>::const_iterator shotEnd,
			 MveDecoder const& mve)
  : shotIt_(shotIt), shotEnd_(shotEnd),
    frameIt_(shotIt_->VGAs.cbegin()),
    audioIt_(shotIt_->AUDIs.cbegin()),
    mve_(mve)
{
}

bool MveDecoder::Movie::hasNext() const {
  if (shotIt_ == shotEnd_)
    return false;

  return true;
}

std::tuple<MveDecoder::Frame const&, MveDecoder::Palette const&> MveDecoder::Movie::decodeNext()
{
  frame_ = mve_.parseVGA(**frameIt_++, &frame_);
  auto& palette = shotIt_->palt;

  if (frameIt_ == shotIt_->VGAs.cend()) {
    ++shotIt_;
    if (shotIt_ != shotEnd_) {
      frameIt_ = shotIt_->VGAs.cbegin();
      audioIt_ = shotIt_->AUDIs.cbegin();
    }
  }
  
  return std::make_tuple(std::ref(frame_), std::ref(palette));
}

bool MveDecoder::Movie::hasAudio() const
{
  return (audioIt_ != shotIt_->AUDIs.cend());
}

MveDecoder::Audio MveDecoder::Movie::getNextAudio()
{
  MveDecoder::Audio ret;
  if (audioIt_ != shotIt_->AUDIs.cend()) {
    ret.reserve((shotIt_->AUDIs.cend()-audioIt_)*2940);

    for(;audioIt_ != shotIt_->AUDIs.cend();++audioIt_)
      std::copy((*audioIt_)->begin(), (*audioIt_)->end(),
		std::back_inserter(ret));
  }

  return ret;
}

unsigned MveDecoder::Movie::getWidth() const
{
  return mve_.width_;
}

unsigned MveDecoder::Movie::getHeight() const
{
  return mve_.height_;
}

MveDecoder::Frame MveDecoder::parseVGA(IffFile::Object const& vga, Frame const* prev) const
{
  // Parse header
  if (vga.getSize() < 8)
    throw FormatException{"VGA chunk smaller than VGA header"};
  
  std::array<uint16_t, 4> segOfs;
  for(unsigned i = 0;i < segOfs.size();++i) {
    segOfs[i] = readU16LE(vga.begin()+i*2);
    if (segOfs[i] >= vga.getSize())
      throw FormatException{"Segment offset exceeds chunk"};
  }

  // Parse segment 1
  if (!segOfs[0] ||
      ((segOfs[1] && (segOfs[1] - segOfs[0]) < 45)) ||
      ((segOfs[2] && (segOfs[2] - segOfs[0]) < 45)) ||
      ((segOfs[3] && (segOfs[3] - segOfs[0]) < 45)) ||
      (vga.getSize() - segOfs[0] < 45))
    throw FormatException{"Segment 1 too small"};

  auto commands = parseHuff_(vga.begin()+segOfs[0], segOfs[1]-segOfs[0]);

  std::vector<uint8_t> pixels;
  if (segOfs[3])
    pixels = parsePixels_(vga.begin()+segOfs[3], vga.getSize()-segOfs[3]);

#ifdef VIDDEBUG_
  printf("%lu commands, %lu pixel data, %u mvecs, %lu size data\n",
	 commands.size(), pixels.size(), (segOfs[2]?(segOfs[3]-segOfs[2]):0),
	 segOfs[2]?(segOfs[2]-segOfs[1]):(segOfs[3]?(segOfs[3]-segOfs[1]):(vga.getSize()-segOfs[1])));
#endif
  
  // Interpret command stream to render frame
  Frame ret;
  ret.reserve(width_*height_);
  
  bool flag = false;
  size_t seg2Idx = 0;
  size_t seg3Idx = 0;
  size_t seg4Idx = 0;
  unsigned size = 0;
  auto it = commands.begin();
  while (it != commands.end()) {
    auto cmd = *it++;
#ifdef CMDDEBUG_
    printf("CMD %u ", cmd);
#endif
    
    switch(cmd) {
    case 0:
      flag = !flag;
      continue;
    case 1:
    case 2:
    case 3:
    case 4:
    case 5:
    case 6:
    case 7:
    case 8:
      size = cmd;
      break;
    case 12:
    case 13:
    case 14:
    case 15:
    case 16:
    case 17:
    case 18:
      size = (cmd - 10);
      break;
    case 9:
    case 19:
      if (!segOfs[1] ||
	  (segOfs[2] && (segOfs[1]+seg2Idx >= segOfs[2])) ||
	  (segOfs[3] && (segOfs[1]+seg2Idx >= segOfs[3])) ||
	  (segOfs[1]+seg2Idx >= vga.getSize()))
	throw FormatException{"Segment 2 overrun"};
      size = vga.begin()[segOfs[1]+seg2Idx++]&0xff;
      break;
    case 10:
    case 20:
      if (!segOfs[1] ||
	  (segOfs[2] && (segOfs[1]+seg2Idx+1 >= segOfs[2])) ||
	  (segOfs[3] && (segOfs[1]+seg2Idx >= segOfs[3])) ||
	  (segOfs[1]+seg2Idx+1 >= vga.getSize()))
	throw FormatException{"Segment 2 overrun"};
      size = readU16BE(vga.begin()+segOfs[1]+seg2Idx);
      seg2Idx += 2;
      break;
    case 11:
    case 21:
      if (!segOfs[1] ||
	  (segOfs[2] && (segOfs[1]+seg2Idx+2 >= segOfs[2])) ||
	  (segOfs[3] && (segOfs[1]+seg2Idx >= segOfs[3])) ||
	  (segOfs[1]+seg2Idx+2 >= vga.getSize()))
	throw FormatException{"Segment 2 overrun"};
      size = readU24BE(vga.begin()+segOfs[1]+seg2Idx);
      seg2Idx += 3;
      break;
    default:
      throw FormatException{"Invalid command"};
    }

    if (cmd < 12) {
      flag = !flag;
      if (flag) {
#ifdef CMDDEBUG_
	printf("PrevCopy: %u\n", size);
#endif
	if (!prev)
	  throw FormatException{"Reference to non-existing prev. frame"};
	if (ret.size()+size > prev->size())
	  throw FormatException{"Copy from prev exceeds frame"};

	std::copy(prev->begin()+ret.size(),
		  prev->begin()+ret.size()+size,
		  std::back_inserter(ret));
      } else {
#ifdef CMDDEBUG_
	printf("DataCopy: %u @%lu\n", size, seg4Idx);
#endif
	if (seg4Idx+size > pixels.size())
	  throw FormatException{"Copy from seg4 exceeds bounds"};

	std::copy(pixels.begin()+seg4Idx,
		  pixels.begin()+seg4Idx+size,
		  std::back_inserter(ret));
	seg4Idx += size;
      }
    } else {
#ifdef CMDDEBUG_
      printf("MvecCopy: %u\n", size);
#endif
      if (!segOfs[2] ||
	  (segOfs[3] && (segOfs[2]+seg3Idx > segOfs[3])) ||
	  (segOfs[2]+seg3Idx > vga.getSize()))
	throw FormatException{"Segment 3 overrun"};
      if (!prev)
	  throw FormatException{"Reference to non-existing prev. frame"};
      
      const uint8_t mpos = vga.begin()[segOfs[2]+seg3Idx++];
      const signed x = sextend(mpos>>4, 4), y = sextend(mpos&0xf,4);
      const signed delta = y*width_+x;
      if ((delta < 0) && (ret.size() < static_cast<size_t>(-delta)))
	throw FormatException{"Motion vector outside frame"};

      const unsigned ref = ret.size()+delta;
#ifdef CMDDEBUG_
      printf("@%lu %.2hhx -> (%d, %d) -> d  %d -> %u\n", seg3Idx-1, mpos, x, y, delta, ref);
#endif

      if (ref+size > prev->size())
	throw FormatException{"Copy from prev exceeds frame"};

      std::copy(prev->begin()+ref,
      		prev->begin()+ref+size,
      		std::back_inserter(ret));
      flag = false;
    }
  }

  if (ret.size() != width_*height_)
    throw FormatException{"Decoded frame dimension mismatch"};

#ifndef NDEBUG
  if (seg4Idx < pixels.size())
    printf("%lu unused pixel data\n", pixels.size()-seg4Idx);
#endif
  
  return ret;
}
/*
#include <SDL2/SDL.h>

struct AudioCBData {
  SDL_AudioSpec spec;
  MveDecoder::Shot const *cur, *next;
  std::vector<IffFile::Object const*>::const_iterator aCur;
  
  size_t pos;
};

void audioCB_(void* userdata, uint8_t *buf, int len)
{
  AudioCBData *cbdata = static_cast<AudioCBData*>(userdata);

  if (!cbdata->cur && cbdata->next) {
    cbdata->cur = cbdata->next;
    cbdata->next = nullptr;
    cbdata->aCur = cbdata->cur->AUDIs.begin();
    cbdata->pos = 0;
  }
  
  while ((len > 0) && cbdata->cur) {
    if (cbdata->pos >= (*cbdata->aCur)->getSize()) {
      cbdata->pos = 0;
      ++cbdata->aCur;
      if (cbdata->aCur == cbdata->cur->AUDIs.end()) {
	cbdata->cur = cbdata->next;
	if (cbdata->next) {
	  cbdata->next = nullptr;
	  cbdata->aCur = cbdata->cur->AUDIs.begin();
	  continue;
	}

	break;
      }
    }
    
    auto& aCur = *cbdata->aCur;
    size_t toCopy = len;
    if (cbdata->pos+toCopy > aCur->getSize())
      toCopy = aCur->getSize()-cbdata->pos;
    std::copy(aCur->begin()+cbdata->pos, aCur->begin()+cbdata->pos+toCopy,
	      buf);
    len -= toCopy;
    buf += toCopy;
    cbdata->pos += toCopy;
  }

  if (len) {
    memset(buf, cbdata->spec.silence, len);
    return;
  }
}

void MveDecoder::play(unsigned branch) const
{
  SDL_Window *window;
  if (SDL_Init(SDL_INIT_VIDEO | SDL_INIT_AUDIO) < 0) {
    printf("Could not init SDL: %s\n", SDL_GetError());
    return;
  }

  window = SDL_CreateWindow("SDL2 Test",
			    SDL_WINDOWPOS_UNDEFINED,
			    SDL_WINDOWPOS_UNDEFINED,
			    width_*2,
			    // 1.2 scale in height to compensate for non-square pixels in height
			    // in original MCGA/VGA 320x200 mode
			    height_*2*1.2,
			    0);
  
  if (!window) {
    printf("Could not create window: %s\n", SDL_GetError());
    SDL_Quit();
    return;
  }

  SDL_Surface *screen = SDL_GetWindowSurface(window);
  if (!screen) {
    printf("Could not get window surface: %s\n", SDL_GetError());
    SDL_DestroyWindow(window);
    SDL_Quit();
    return;
  }

  

  SDL_Palette *palette = SDL_AllocPalette(256);
  if(!palette) {
    printf("Could not create palette: %s\n", SDL_GetError());
    SDL_DestroyWindow(window);
    SDL_Quit();
    return;
  }
  
  auto shot = branches_.at(branch).begin();
  auto curShot = shot++;
  
  std::array<SDL_Color, 256> colors;
  for(unsigned i = 0;i < 256;++i) {
    colors[i].r = curShot->palt[i*3];
    colors[i].g = curShot->palt[i*3+1];
    colors[i].b = curShot->palt[i*3+2];
    colors[i].a = 255;
  }
  
  if (SDL_SetPaletteColors(palette, colors.data(), 0, 256) != 0) {
    printf("Could not set palette: %s\n", SDL_GetError());
    SDL_FreePalette(palette);
    SDL_DestroyWindow(window);
    SDL_Quit();
    return;
  }

  AudioCBData audioCBData;
  if (curShot->AUDIs.size() > 0) {
    audioCBData.cur = &*curShot;
    audioCBData.aCur = curShot->AUDIs.cbegin();
  } else {
    audioCBData.cur = nullptr;
  }
  audioCBData.next = nullptr;
  audioCBData.pos = 0;
  
  SDL_AudioSpec want;
  want.freq = 22050;
  want.format = AUDIO_S16LSB;
  want.channels = 1;
  want.samples = 4096;
  want.callback = &audioCB_;
  want.userdata = &audioCBData;

  auto audioDev = SDL_OpenAudioDevice(nullptr, false, &want, &audioCBData.spec, 0);
  if (!audioDev) {
    printf("Failed to open audio: %s\n", SDL_GetError());
    SDL_FreePalette(palette);
    SDL_DestroyWindow(window);
    SDL_Quit();
    return;
  }
  
  auto vga = curShot->VGAs.begin();
  uint32_t last = 0;
  const double frameTime = 1000/15.0;
  double nextFT = frameTime;
  Frame frame = parseVGA(**vga++, nullptr);  

  decltype(last) minFT = std::numeric_limits<decltype(last)>::max(), maxFT = 0;
  
  bool close = false;

  SDL_PauseAudioDevice(audioDev, 0);
  while (!close) {
    SDL_Event event;
    while (SDL_PollEvent(&event)) {
      switch(event.type) {
      case SDL_KEYDOWN:
	close = true;
	break;
      case SDL_WINDOWEVENT:
	if (event.window.event == SDL_WINDOWEVENT_CLOSE)
	  close = true;
	break;
      }
    }
    
    auto now = SDL_GetTicks();
    if (now-last >= nextFT) {
      if (last) {
        minFT = std::min(minFT, (now-last));
	maxFT = std::max(maxFT, (now-last));
	
	if ((now-last)-nextFT <= frameTime)
	  nextFT = frameTime-((now-last)-nextFT);
	else {
	  printf("NYI: Frame dropping\n");
	  nextFT = frameTime;
	}
      }
      
      SDL_Surface *vidSurf = SDL_CreateRGBSurfaceFrom((void*)frame.pixels.data(), width_, height_, 8, width_, 0, 0, 0, 0);
      if (!vidSurf) {
	printf("Could not get video surface: %s\n", SDL_GetError());
	SDL_CloseAudioDevice(audioDev);
	SDL_FreePalette(palette);
	SDL_DestroyWindow(window);
	SDL_Quit();
	return;
      }

      if (SDL_SetSurfacePalette(vidSurf, palette) != 0) {
	printf("Could not set surface palette: %s\n", SDL_GetError());
	SDL_FreeSurface(vidSurf);
	SDL_CloseAudioDevice(audioDev);
	SDL_FreePalette(palette);
	SDL_DestroyWindow(window);
	SDL_Quit();
	return;
      }

      SDL_Surface *vidSurfRGB = SDL_ConvertSurface(vidSurf, screen->format, 0);
      if (!vidSurfRGB) {
	printf("Could not convert video surface: %s\n", SDL_GetError());
	SDL_FreeSurface(vidSurf);
	SDL_CloseAudioDevice(audioDev);
	SDL_FreePalette(palette);
	SDL_DestroyWindow(window);
	SDL_Quit();
	return;
      }
    
      if (SDL_BlitScaled(vidSurfRGB, nullptr, screen, nullptr) != 0) {
	printf("Could not blit video surface: %s\n", SDL_GetError());
	SDL_FreeSurface(vidSurfRGB);
	SDL_FreeSurface(vidSurf);
	SDL_CloseAudioDevice(audioDev);
	SDL_FreePalette(palette);
	SDL_DestroyWindow(window);
	SDL_Quit();
	return;
      }

      SDL_UpdateWindowSurface(window);
      last = now;
      SDL_FreeSurface(vidSurfRGB);
      SDL_FreeSurface(vidSurf);

      // Decode next frame
      if (vga == curShot->VGAs.end()) {
	if (shot == branches_.at(branch).end())
	  break;
	else {
	  curShot = shot++;
	  vga = curShot->VGAs.begin();
	
	  std::array<SDL_Color, 256> colors;
	  for(unsigned i = 0;i < 256;++i) {
	    colors[i].r = curShot->palt[i*3];
	    colors[i].g = curShot->palt[i*3+1];
	    colors[i].b = curShot->palt[i*3+2];
	    colors[i].a = 255;
	  }
  
	  if (SDL_SetPaletteColors(palette, colors.data(), 0, 256) != 0) {
	    printf("Could not set palette: %s\n", SDL_GetError());
	    SDL_CloseAudioDevice(audioDev);
	    SDL_FreePalette(palette);
	    SDL_DestroyWindow(window);
	    SDL_Quit();
	    return;
	  }
	  SDL_LockAudioDevice(audioDev);
	  audioCBData.next = &*curShot;
	  SDL_UnlockAudioDevice(audioDev);
	}
      }
      frame = parseVGA(**vga++, &frame);  
    } else if ((nextFT-(now-last)) >= 10)
      SDL_Delay(nextFT-(now-last));
  }
  
  SDL_PauseAudioDevice(audioDev, 1);
  SDL_CloseAudioDevice(audioDev);
  SDL_FreePalette(palette);
  SDL_DestroyWindow(window);
  SDL_Quit();

  printf("Frame times: [%u, %u]\n",
	 minFT, maxFT);
}
*/