/** @file
  Copyright (C) 2019, Goldfish64. All rights reserved.

  This program and the accompanying materials
  are licensed and made available under the terms and conditions of the BSD License
  which accompanies this distribution.  The full text of the license may be found at
  http://opensource.org/licenses/bsd-license.php

  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
**/

#include <Base.h>

#include <Library/BaseLib.h>
#include <Library/DebugLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/OcAppleDiskImageLib.h>
#include <Library/OcXmlLib.h>

#include "OcAppleDiskImageLibInternal.h"

STATIC
BOOLEAN
InternalFindPlistDictChild (
  IN  XML_NODE  *Node,
  IN  CHAR8     *KeyName,
  OUT XML_NODE  **Key,
  OUT XML_NODE  **Value
  )
{
  UINT32      ChildCount;
  XML_NODE    *ChildValue;
  XML_NODE    *ChildKey;
  CONST CHAR8 *ChildKeyName;
  UINT32      Index;

  ASSERT (Node != NULL);
  ASSERT (KeyName != NULL);
  ASSERT (Key != NULL);
  ASSERT (Value != NULL);

  ChildCount = PlistDictChildren (Node);
  for (Index = 0; Index < ChildCount; ++Index) {
    ChildKey     = PlistDictChild (Node, Index, &ChildValue);
    ChildKeyName = PlistKeyValue (ChildKey);
    if (ChildKeyName == NULL) {
      break;
    }

    if (AsciiStrCmp (ChildKeyName, KeyName) == 0) {
      *Key   = ChildKey;
      *Value = ChildValue;
      return TRUE;
    }
  }

  return FALSE;
}

STATIC
VOID
InternalSwapBlockData (
  IN OUT APPLE_DISK_IMAGE_BLOCK_DATA  *BlockData
  )
{
  APPLE_DISK_IMAGE_CHUNK *Chunk;
  UINT32                 Index;

  BlockData->Version          = SwapBytes32 (BlockData->Version);
  BlockData->SectorNumber     = SwapBytes64 (BlockData->SectorNumber);
  BlockData->SectorCount      = SwapBytes64 (BlockData->SectorCount);
  BlockData->DataOffset       = SwapBytes64 (BlockData->DataOffset);
  BlockData->BuffersNeeded    = SwapBytes32 (BlockData->BuffersNeeded);
  BlockData->BlockDescriptors = SwapBytes32 (BlockData->BlockDescriptors);
  BlockData->Checksum.Type    = SwapBytes32 (BlockData->Checksum.Type);
  BlockData->Checksum.Size    = SwapBytes32 (BlockData->Checksum.Size);

  for (Index = 0; Index < APPLE_DISK_IMAGE_CHECKSUM_SIZE; ++Index) {
    BlockData->Checksum.Data[Index] = SwapBytes32 (
                                        BlockData->Checksum.Data[Index]
                                        );
  }

  BlockData->ChunkCount = SwapBytes32 (BlockData->ChunkCount);

  for (Index = 0; Index < BlockData->ChunkCount; ++Index) {
    Chunk = &BlockData->Chunks[Index];

    Chunk->Type             = SwapBytes32 (Chunk->Type);
    Chunk->Comment          = SwapBytes32 (Chunk->Comment);
    Chunk->SectorNumber     = SwapBytes64 (Chunk->SectorNumber);
    Chunk->SectorCount      = SwapBytes64 (Chunk->SectorCount);
    Chunk->CompressedOffset = SwapBytes64 (Chunk->CompressedOffset);
    Chunk->CompressedLength = SwapBytes64 (Chunk->CompressedLength);
  }
}

BOOLEAN
InternalParsePlist (
  IN  VOID                         *Buffer,
  IN  UINT32                       XmlOffset,
  IN  UINT32                       XmlLength,
  OUT UINT32                       *BlockCount,
  OUT APPLE_DISK_IMAGE_BLOCK_DATA  ***Blocks
  )
{
  BOOLEAN                     Result;

  CHAR8                       *XmlPlistBuffer;
  XML_DOCUMENT                *XmlPlistDoc;
  XML_NODE                    *NodeRoot;
  XML_NODE                    *NodeResourceForkKey;
  XML_NODE                    *NodeResourceForkValue;
  XML_NODE                    *NodeBlockListKey;
  XML_NODE                    *NodeBlockListValue;

  XML_NODE                    *NodeBlockDict;
  XML_NODE                    *BlockDictChildKey;
  XML_NODE                    *BlockDictChildValue;
  UINT32                      BlockDictChildDataSize;

  UINT32                      NumDmgBlocks;
  APPLE_DISK_IMAGE_BLOCK_DATA **DmgBlocks;
  APPLE_DISK_IMAGE_BLOCK_DATA *Block;

  UINT32 Index;

  ASSERT (Buffer != NULL);
  ASSERT (XmlLength > 0);
  ASSERT (BlockCount != NULL);
  ASSERT (Blocks != NULL);

  DmgBlocks = NULL;

  XmlPlistDoc = NULL;

  XmlPlistBuffer = AllocateCopyPool (XmlLength, ((UINT8 *)Buffer + XmlOffset));
  if (XmlPlistBuffer == NULL) {
    Result = FALSE;
    goto DONE_ERROR;
  }

  XmlPlistDoc = XmlDocumentParse (XmlPlistBuffer, XmlLength, FALSE);
  if (XmlPlistDoc == NULL) {
    Result = FALSE;
    goto DONE_ERROR;
  }

  NodeRoot = PlistDocumentRoot (XmlPlistDoc);
  if (NodeRoot == NULL) {
    Result = FALSE;
    goto DONE_ERROR;
  }

  Result = InternalFindPlistDictChild (
             NodeRoot,
             DMG_PLIST_RESOURCE_FORK_KEY,
             &NodeResourceForkKey,
             &NodeResourceForkValue
             );
  if (!Result) {
    goto DONE_ERROR;
  }

  Result = InternalFindPlistDictChild (
             NodeResourceForkValue,
             DMG_PLIST_BLOCK_LIST_KEY,
             &NodeBlockListKey,
             &NodeBlockListValue
             );
  if (!Result) {
    goto DONE_ERROR;
  }

  NumDmgBlocks = XmlNodeChildren (NodeBlockListValue);
  if (NumDmgBlocks == 0) {
    Result = FALSE;
    goto DONE_ERROR;
  }

  DmgBlocks = AllocatePool (NumDmgBlocks * sizeof (*DmgBlocks));
  if (DmgBlocks == NULL) {
    Result = FALSE;
    goto DONE_ERROR;
  }

  for (Index = 0; Index < NumDmgBlocks; ++Index) {
    NodeBlockDict = XmlNodeChild (NodeBlockListValue, Index);

    Result = InternalFindPlistDictChild (
               NodeBlockDict,
               DMG_PLIST_DATA,
               &BlockDictChildKey,
               &BlockDictChildValue
               );
    if (!Result) {
      goto DONE_ERROR;
    }

    Result = PlistDataSize (BlockDictChildValue, &BlockDictChildDataSize);
    if (!Result) {
      goto DONE_ERROR;
    }

    Block = AllocatePool (BlockDictChildDataSize);
    if (Block == NULL) {
      FreePool (Block);
      Result = FALSE;
      goto DONE_ERROR;
    }

    Result = PlistDataValue (
               BlockDictChildValue,
               (UINT8 *)Block,
               &BlockDictChildDataSize
               );
    if (!Result) {
      FreePool (Block);
      goto DONE_ERROR;
    }

    InternalSwapBlockData (Block);
    DmgBlocks[Index] = Block;
  }

  *BlockCount = NumDmgBlocks;
  *Blocks     = DmgBlocks;
  Result      = TRUE;

DONE_ERROR:
  if (!Result && (DmgBlocks != NULL)) {
    while ((Index--) != 0) {
      FreePool (DmgBlocks[Index]);
    }

    FreePool (DmgBlocks);
  }

  if (XmlPlistDoc != NULL) {
    XmlDocumentFree (XmlPlistDoc);
  }

  if (XmlPlistBuffer != NULL) {
    FreePool (XmlPlistBuffer);
  }

  return Result;
}

BOOLEAN
InternalGetBlockChunk (
  IN  OC_APPLE_DISK_IMAGE_CONTEXT  *Context,
  IN  UINT64                       Lba,
  OUT APPLE_DISK_IMAGE_BLOCK_DATA  **Data,
  OUT APPLE_DISK_IMAGE_CHUNK       **Chunk
  )
{
  UINT32                      BlockIndex;
  UINT32                      ChunkIndex;
  APPLE_DISK_IMAGE_BLOCK_DATA *BlockData;
  APPLE_DISK_IMAGE_CHUNK      *BlockChunk;

  for (BlockIndex = 0; BlockIndex < Context->BlockCount; ++BlockIndex) {
    BlockData = Context->Blocks[BlockIndex];

    if ((Lba >= BlockData->SectorNumber)
     && (Lba < (BlockData->SectorNumber + BlockData->SectorCount))) {
      for (ChunkIndex = 0; ChunkIndex < BlockData->ChunkCount; ++ChunkIndex) {
        BlockChunk = &BlockData->Chunks[ChunkIndex];

        if ((Lba >= DMG_SECTOR_START_ABS (BlockData, BlockChunk))
         && (Lba < (DMG_SECTOR_START_ABS (BlockData, BlockChunk) + BlockChunk->SectorCount))) {
          *Data  = BlockData;
          *Chunk = BlockChunk;
          return TRUE;
        }
      }
    }
  }

  return FALSE;
}