extern crate lzw;

use std::cell::RefCell;

use self::lzw::BitReader;

static EGA_PALETTE: [[u8; 4]; 16] = [[0x00, 0x00, 0x00, 0xFF],
                                     [0x00, 0x00, 0xAA, 0xFF],
                                     [0x00, 0xAA, 0x00, 0xFF],
                                     [0x00, 0xAA, 0xAA, 0xFF],
                                     [0xAA, 0x00, 0x00, 0xFF],
                                     [0xAA, 0x00, 0xAA, 0xFF],
                                     [0xAA, 0x55, 0x00, 0xFF],
                                     [0xAA, 0xAA, 0xAA, 0xFF],
                                     [0x55, 0x55, 0x55, 0xFF],
                                     [0x55, 0x55, 0xFF, 0xFF],
                                     [0x55, 0xFF, 0x55, 0xFF],
                                     [0x55, 0xFF, 0xFF, 0xFF],
                                     [0xFF, 0x55, 0x55, 0xFF],
                                     [0xFF, 0x55, 0xFF, 0xFF],
                                     [0xFF, 0xFF, 0x55, 0xFF],
                                     [0xFF, 0xFF, 0xFF, 0xFF]];

pub enum Compression {
  Uncompressed,
  Rle,
  Lzw
}

pub enum Tiling {
  Untiled,
  Tiled(u16)
}

pub struct EgaPage {
  pub data: Vec<u8>,
  pub dim: usize,
}

impl EgaPage {
  pub fn iter<'a>(&'a self) -> impl Iterator<Item=&'a [u8]> {
    self.data.chunks(4 * self.dim * self.dim)
  }

  pub fn mipmap(&self, scale_levels: u8) -> Mipmap {
    let scale = 1 << scale_levels;
    let scaled: Vec<Vec<u8>> = self.iter().map(|tile| {
      let mut tilevec = Vec::new();
      for row in tile.chunks(4 * self.dim) {
        for _ in 0..scale {
          for px in row.chunks(4) {
            for subpx in px.iter().cycle().take(4 * scale) {
              tilevec.push(*subpx);
            }
          }
        }
      }
      tilevec
    }).collect();

    // dummy buffers just to provide proper sizing on texture for mips
    let dim = self.dim << scale_levels;
    let mut dummies = Vec::new();

    let mut level = 1;
    loop {
      let mipdim = dim >> level;
      if mipdim == 0 { break; }
      dummies.push(vec![0; 4 * mipdim * mipdim]);
      level += 1;
    }

    Mipmap {
      len: self.data.len() / (4 * self.dim * self.dim),
      dim: dim,
      backing: scaled,
      dummies: dummies,
    }
  }
}

pub struct Mipmap {
  pub len: usize,
  pub dim: usize,
  backing: Vec<Vec<u8>>,
  dummies: Vec<Vec<u8>>,
}

impl Mipmap {
  pub fn slices(&self) -> Vec<&[u8]> {
    let mut slices: Vec<&[u8]> = Vec::with_capacity((1 + self.dummies.len()) * self.len);
    for tile in self.backing.iter() {
      slices.push(tile);
      for dummy in self.dummies.iter() {
        slices.push(&dummy);
      }
    }
    slices
  }
}

pub fn decode<'a>(buf: &[u8], compression: Compression, tiling: Tiling)
                  -> EgaPage {
  let out: Vec<u8>;

  out = match compression {
    Compression::Uncompressed => decode_uncompressed(buf),
    Compression::Rle => {
      decode_uncompressed(&decode_rle(buf))
    },
    Compression::Lzw => {
      let mut decoder = U4Lzw::new();
      decode_uncompressed(&decoder.decode(buf))
    },
  };
  let dim = match tiling {
    Tiling::Tiled(tiledim) => tiledim as usize,
    Tiling::Untiled => out.len()
  };
  EgaPage { data: out, dim: dim}
}

fn decode_uncompressed(buf: &[u8]) -> Vec<u8> {
  buf.iter()
     .flat_map(|tile_byte| {
         EGA_PALETTE[(tile_byte >> 4u8 & 0xF) as usize]
         .into_iter()
         .chain(EGA_PALETTE[(tile_byte & 0xF) as usize]
                .into_iter())
     })
    .map(|x| *x)
    .collect()
}

fn decode_rle(bytes: &[u8]) -> Vec<u8> {
  let mut out = Vec::new();
  let mut iter = bytes.iter();

  while let Some(cmd) = iter.next() {
    match cmd {
      &0x02 => {
        let len = *iter.next().expect("rle missing run length") as usize;
        out.extend(::std::iter::repeat(*iter.next().expect("rle missing run value")).take(len));
      },
      val => out.push(*val),
    }
  }
  out
}

struct U4Lzw {
  table: Vec<RefCell<Vec<u8>>>,
  load: u16,
}

const U4_CODE_SIZE: u8 = 12;
impl U4Lzw {
  fn new() -> U4Lzw {
    let mut out = U4Lzw {
      table: vec![RefCell::new(Vec::new()); 1 << U4_CODE_SIZE],
      load: 0
    };
    for (i, entry) in out.table.iter_mut().take(256).enumerate() {
      entry.borrow_mut().push(i as u8)
    }
    out
  }

  fn reset(&mut self) {
    for i in 256..(1 << U4_CODE_SIZE) {
      self.table[i].borrow_mut().clear();
    }
    self.load = 0;
  }

  fn decode(&mut self, bytes: &[u8]) -> Vec<u8> {
    let mut out = Vec::new();
    let mut iter = CodeIter::new(bytes);

    let mut prev = iter.next().expect("empty lzw");
    let mut ch = prev as u8;
    out.push(ch);

    for code in iter {
      {
        let mut new;
        let mut seq = &*self.table[code as usize].borrow();
        if seq.is_empty() {
          new = self.table[prev as usize].borrow().clone();
          new.push(ch);
          seq = &new;
        };
        out.extend(seq.iter());
        ch = seq[0];
      }
      let hash = self.hash(ch, prev);
      {
        let old = self.table[prev as usize].borrow();
        let mut entry = self.table[hash].borrow_mut();
        entry.extend(old.iter());
        entry.push(ch);
      }
      prev = code;

      self.load += 1;
      if self.load > 0xccd { self.reset(); }
    }
    out
  }

  fn hash(&self, root: u8, code: u16) -> usize {
    let hash = Self::hash_primary(root, code);
    if self.table[hash].borrow().is_empty() { return hash; }

    let mut hash = Self::hash_secondary(root, code);
    while !self.table[hash].borrow().is_empty() {
      hash = (hash + 509) & 0xfff
    }

    hash
  }

  fn hash_primary(root: u8, code: u16) -> usize {
    ((root as (u16) << 4) ^ code) as usize
  }

  fn hash_secondary(root: u8, code: u16) -> usize {
    let base = (((root as (u16) << 1) + code) | 0x800) as u32;
    let squared = base * base;
    ((squared & 0x0003ffc0) >> 6) as usize
  }
}

struct CodeIter<'a> {
  pos: usize,
  bytes: &'a [u8],
  reader: lzw::MsbReader,
}

impl<'a> CodeIter<'a> {
  fn new(bytes: &[u8]) -> CodeIter {
    CodeIter { pos: 0, bytes: bytes, reader: lzw::MsbReader::new() }
  }
}

impl<'a> Iterator for CodeIter<'a> {
  type Item = u16;
  fn next(&mut self) -> Option<u16> {
    match self.reader.read_bits(&self.bytes[self.pos..], U4_CODE_SIZE) {
      lzw::Bits::Some(used, code) => {
        self.pos += used;
        Some(code)
      },
      lzw::Bits::None(used) => {
        self.pos += used;
        None
      }
    }
  }
}