using System;
using System.Text;
using System.IO;
using System.Linq;
using System.IO.Compression;
using System.Collections.Generic;
using System.Diagnostics;
using System.Threading.Tasks;

public class MMArchiveHelper
{
    public static void CheckMagic(uint magic, uint required)
    {
        if (magic != required)
        {
            throw new Exception(string.Format("Invalid Magic {0:08X} (expected {1:08X})", magic, required));
        }
    }

    public static string ReadASCII(Stream stream)
    {
        var builder = new StringBuilder();

        for (int i; (i = stream.ReadByte()) > 0;)
        {
            builder.Append((char) i);
        }

        return builder.ToString();
    }

    public class ArchiveEntry
    {
        public readonly string Name;
        public readonly long Size;
        public readonly long RawSize;
        public readonly long Position;

        public ArchiveEntry(string name, long size, long rawSize, long fileOffset)
        {
            Name     = name;
            Size     = size;
            RawSize  = rawSize;
            Position = fileOffset;
        }

        public byte[] GetContents(Stream input)
        {
            input.Seek(Position, SeekOrigin.Begin);

            if (Size != RawSize)
            {
                input = new DeflateStream(input, CompressionMode.Decompress, true);
            }

            var buffer = new byte[Size];

            input.Read(buffer, 0, buffer.Length);

            return buffer;
        }

        public override string ToString()
        {
            return Name;
        }
    }

    public class VirtualFileInode
    {
        public readonly string Name;
        public readonly uint DataOffset;
        public readonly uint Size;
        public readonly bool IsDirectory;

        public VirtualFileInode(BinaryReader reader, Stream nameStream)
        {
            DataOffset = reader.ReadUInt32();

            var field4 = reader.ReadUInt32();
            var field8 = reader.ReadUInt32();

            Size        = (field4 & 0x7FFFFF);
            IsDirectory = (field8 & 1) != 0;

            var nameOffset      = (field8 >> 14) & 0x3FFFF;
            var extensionOffset = (field4 >> 23) & 0x1FF;
            var nameInteger     = (field8 >>  1) & 0x1FFF;

            nameStream.Seek(nameOffset, SeekOrigin.Begin);

            Name = ReadASCII(nameStream).Replace("\x01", nameInteger.ToString());

            if (extensionOffset != 0)
            {
                nameStream.Seek(extensionOffset, SeekOrigin.Begin);

                Name += "." + ReadASCII(nameStream);
            }
        }

        public ArchiveEntry GetArchiveEntry(string parent = "")
        {
            if (!IsDirectory)
            {
                return new ArchiveEntry(parent + Name, Size, Size, DataOffset);
            }

            return null;
        }

        public IEnumerable<VirtualFileInode> GetChildren(VirtualFileInode[] nodes)
        {
            if (IsDirectory)
            {
                for (var i = DataOffset; i < DataOffset + Size; ++i)
                {
                    yield return nodes[i];
                }
            }
        }

        public IEnumerable<ArchiveEntry> GetFiles(VirtualFileInode[] nodes, string parent = "")
        {
            if (IsDirectory)
            {
                var name = parent + Name + "\\";

                foreach (var node in GetChildren(nodes))
                {
                    if (node.IsDirectory)
                    {
                        foreach (var file in node.GetFiles(nodes, name))
                        {
                            yield return file;
                        }
                    }
                    else
                    {
                        yield return node.GetArchiveEntry(name);
                    }
                }
            }
        }

        public override string ToString()
        {
            return Name;
        }
    }

    public static IEnumerable<ArchiveEntry> ReadDaveArchive(BinaryReader reader)
    {
        reader.BaseStream.Seek(0, SeekOrigin.Begin);

        CheckMagic(reader.ReadUInt32(), 0x45564144);

        var headerCount = reader.ReadUInt32();
        var namesOffset = reader.ReadUInt32();
        var namesSize   = reader.ReadUInt32();

        for (var i = 0; i < headerCount; ++i)
        {
            reader.BaseStream.Seek(2048 + (i * 16), SeekOrigin.Begin);

            var nameOffset = reader.ReadUInt32();
            var dataOffset = reader.ReadUInt32();

            var uncompressedSize = reader.ReadUInt32();
            var compressedSize   = reader.ReadUInt32();

            reader.BaseStream.Seek(2048 + namesOffset + nameOffset, SeekOrigin.Begin);

            var name = ReadASCII(reader.BaseStream);

            yield return new ArchiveEntry(name, uncompressedSize, compressedSize, dataOffset);
        }
    }

    public static IEnumerable<ArchiveEntry> ReadPKArchive(BinaryReader reader)
    {
        reader.BaseStream.Seek(-22, SeekOrigin.End);

        CheckMagic(reader.ReadUInt32(), 0x06054B50); // ZIPENDLOCATOR

        var diskNumber      = reader.ReadUInt16();
        var startDiskNumber = reader.ReadUInt16();

        if (diskNumber != startDiskNumber)
        {
            throw new Exception("Incomplete Archive");
        }

        var fileCount        = reader.ReadUInt16();
        var filesInDirectory = reader.ReadUInt16();

        var directorySize   = reader.ReadUInt32();
        var directoryOffset = reader.ReadUInt32();

        var fileCommentLength = reader.ReadUInt16();

        var currentOffset = directoryOffset;

        while (true)
        {
            reader.BaseStream.Seek(currentOffset, SeekOrigin.Begin);

            if (reader.ReadUInt32() != 0x02014B50) // ZIPDIRENTRY
            {
                break;
            }

            var versionMadeBy    = reader.ReadUInt16();
            var versionToExtract = reader.ReadUInt16();
            var flags            = reader.ReadUInt16();

            var compressionMethod = reader.ReadUInt16();

            if (compressionMethod != 0 && compressionMethod != 8)
            {
                throw new Exception("Invalid compression method " + compressionMethod);
            }

            var fileTime = reader.ReadUInt16();
            var fileDate = reader.ReadUInt16();
            var crc      = reader.ReadUInt32();

            var compressedSize   = reader.ReadUInt32();
            var uncompressedSize = reader.ReadUInt32();

            var nameLength    = reader.ReadUInt16();
            var extraLength   = reader.ReadUInt16();
            var commentLength = reader.ReadUInt16();

            var diskNumberStart = reader.ReadUInt16();

            var internalAttributes = reader.ReadUInt16();
            var externalAttributes = reader.ReadUInt32();

            var dataOffset = reader.ReadUInt32(); // ZIPFILERECORD

            var name = new string(reader.ReadChars(nameLength));

            currentOffset = (uint) reader.BaseStream.Position + extraLength + commentLength;

            yield return new ArchiveEntry(name, uncompressedSize, compressedSize, dataOffset + nameLength + 30);
        }
    }

    public static IEnumerable<ArchiveEntry> ReadARESArchive(BinaryReader reader)
    {
        reader.BaseStream.Seek(0, SeekOrigin.Begin);

        CheckMagic(reader.ReadUInt32(), 0x53455241);

        var headerCount    = reader.ReadUInt32();
        var directoryCount = reader.ReadUInt32();
        var namesSize      = reader.ReadUInt32();
        var namesOffset    = 16 + headerCount * 12;

        reader.BaseStream.Seek(namesOffset, SeekOrigin.Begin);

        var namesStream = new MemoryStream(reader.ReadBytes((int) namesSize));

        var nodes = new VirtualFileInode[headerCount];

        for (var i = 0; i < headerCount; ++i)
        {
            reader.BaseStream.Seek(16 + (i * 12), SeekOrigin.Begin);

            nodes[i] = new VirtualFileInode(reader, namesStream);
        }

        for (var i = 0; i < directoryCount; ++i)
        {
            foreach (var file in nodes[i].GetFiles(nodes, ""))
            {
                yield return file;
            }
        }
    }

    public static IEnumerable<ArchiveEntry> GetArchiveEntries(Stream archive)
    {
        var binaryReader = new BinaryReader(archive);

        binaryReader.BaseStream.Seek(0, SeekOrigin.Begin);

        if (binaryReader.ReadUInt32() == 0x45564144) // DAVE
        {
            return ReadDaveArchive(binaryReader);
        }

        binaryReader.BaseStream.Seek(0, SeekOrigin.Begin);

        if (binaryReader.ReadUInt32() == 0x53455241) // ARES
        {
            return ReadARESArchive(binaryReader);
        }

        binaryReader.BaseStream.Seek(-22, SeekOrigin.End);

        if (binaryReader.ReadUInt32() == 0x06054B50) // ZIPENDLOCATOR
        {
            return ReadPKArchive(binaryReader);
        }

        throw new ArgumentException("Invalid archive stream");
    }
}

namespace DaveReader
{
    class MMArchiveReader
    {
        static string ConsoleClearLine()
        {
            return string.Format("\r{0}\r", new string(' ', Console.WindowWidth - 1));
        }

        static void OverwriteConsole(string format, params object[ ] args)
        {
            Console.Write(ConsoleClearLine() + format, args);
        }

        static void OverwriteConsoleLine(string format, params object[ ] args)
        {
            Console.WriteLine(ConsoleClearLine() + format, args);
        }

        static void ExtractArchive(string path, bool merge)
        {
            using (var fileStream = Stream.Synchronized(File.OpenRead(path)))
            {
                var stopwatch = Stopwatch.StartNew();

                var tasks = MMArchiveHelper.GetArchiveEntries(fileStream).Select((entry) =>
                {
                    try
                    {
                        return new Task(() =>
                        {
                            var contents = entry.GetContents(fileStream);

                            if (contents.Length > 0)
                            {
                                var outputName = merge
                                    ? Path.Combine(Environment.CurrentDirectory, "Extracted", entry.Name)
                                    : Path.Combine(Environment.CurrentDirectory, "Extracted", Path.GetFileNameWithoutExtension(path), entry.Name);

                                Directory.CreateDirectory(Path.GetDirectoryName(outputName));

                                using (var output = File.OpenWrite(outputName))
                                {
                                    output.Write(contents, 0, contents.Length);
                                }
                            }
                        });
                    }
                    catch (Exception e)
                    {
                        Console.WriteLine("Failed to extract {0} - {1}", entry.Name, e);
                    }

                    return null;
                }).ToArray();

                foreach (var task in tasks)
                {
                    task.Start();
                }

                var endTask = Task.WhenAll(tasks);

                while (!endTask.Wait(100))
                {
                    var completed = tasks.Where(x => x.IsCompleted).ToArray();

                    OverwriteConsole("{0} : {1} / {2}", path, completed.Length, tasks.Length);
                }

                OverwriteConsoleLine("{0} : {1} files in {2} ms", path, tasks.Length, stopwatch.ElapsedMilliseconds);
            }
        }

        static string GetArgOrReadLine(string[] args, uint index)
        {
            return args.Length > index ? args[index] : Console.ReadLine();
        }

        static string GetArgOfDefault(string[] args, uint index, string defaultValue)
        {
            return args.Length > index ? args[index] : defaultValue;
        }

        static void Main(string[] args)
        {
            Console.Clear();
            Console.ForegroundColor = ConsoleColor.Green;

            var path  = GetArgOrReadLine(args, 0);
            var merge = bool.Parse(GetArgOfDefault(args, 1, "false"));

            if (File.Exists(path))
            {
                ExtractArchive(path, merge);
            }
            else if (Directory.Exists(path))
            {
                foreach (var file in Directory.GetFiles(path, "*.ar"))
                {
                    ExtractArchive(file, merge);
                }
            }
            else
            {
                Console.WriteLine("Invalid file/path: {0}", path);
            }

            Console.WriteLine("Finished");

            Console.ReadKey(true);
        }
    }
}