using System;
using System.Runtime.InteropServices;
using System.Collections.Generic;
using UdonSharp;
using UnityEngine;
using VRC.SDK3.Rendering;
using VRC.Udon.Common.Interfaces;

public class DataDecoder : UdonSharpBehaviour
{
  public RenderTexture sourceTexture;
  public MeshRenderer target;

  private int tileSize = 8;
  // Minimum size (in pixels) of a tile. This is shared with our tixl operator.
  private const int kMinTileSize = 4;
  private const int kMaxTileSize = 128;
  private Color32[] pixelData;
  private bool hasData = false;
  private int readWidth;
  private int readHeight;

  // The wall time at which we last saw a sync event
  private float wallSyncTime;
  // The logical time corresponding to the last sync event
  private float logicalSyncTimeMs;
  // The rate at which logical time passes every second.
  private float logicalTimeFactor;

  // Top-level data types.
  private const int kT_TimeSyncData = 0;

  void Start() {}

  void Update()
  {
    if (sourceTexture == null) return;

    // Request the rectangular region which leaves either a right-justified
    // square, or a bottom-justified square.
    int requestWidth;
    int requestHeight;
    if (sourceTexture.width < sourceTexture.height) {
      requestWidth = sourceTexture.width;
      requestHeight = sourceTexture.height - sourceTexture.width;
    } else {
      requestHeight = sourceTexture.height;
      requestWidth = sourceTexture.width - sourceTexture.height;
    }
    int pixelCount = requestWidth * requestHeight;

    if (pixelCount <= 0) return;

    if (pixelData == null || pixelCount != pixelData.Length)
    {
      pixelData = new Color32[pixelCount];
      hasData = false;
    }

    readWidth = requestWidth;
    readHeight = requestHeight;

    VRCAsyncGPUReadback.Request(sourceTexture,
        0,
        0, readWidth,
        0, readHeight,
        0, 1,
        (IUdonEventReceiver)this);

    if (hasData)
    {
      ProcessTiles();
      hasData = false;
    }

    if (wallSyncTime != null) {
      float logicalTime = logicalSyncTimeMs * 0.001f +
        logicalTimeFactor * (Time.time - wallSyncTime);
      if (target != null) {
        target.material.SetFloat("_Logical_Time", logicalTime);
      }
    }
  }

  public override void OnAsyncGpuReadbackComplete(VRCAsyncGPUReadbackRequest request)
  {
    if (request.hasError) return;

    if (pixelData != null && request.TryGetData(pixelData))
    {
      hasData = true;
    }
  }

  // Byte parsing logic
  private bool HasBytesLeft(byte[] b, int offset, int bytesLeft) {
    return b.Length - offset >= bytesLeft;
  }

  private int GetInt(ref byte[] b, ref int offset) {
    int ret = BitConverter.ToInt32(b, offset);
    offset += 4;
    return ret;
  }

  private float GetFloat(ref byte[] b, ref int offset) {
    float ret = BitConverter.ToSingle(b, offset);
    offset += 4;
    return ret;
  }

  // Convert a two's complement number in the lower 4 bits of a byte to a
  // gray code in the lower 4 bits.
  private byte ToGrayNibble(byte twosComplementNibble)
  {
    int lowerN = twosComplementNibble & 0x0F;
    int lowerG = lowerN ^ (lowerN >> 1);
    return (byte)((twosComplementNibble & 0xF0) | lowerG);
  }

  // Convert a gray code in the lower 4 bits of a byte to a two's complement
  // number in the lower 4 bits.
  private byte ToTwosComplementNibble(byte grayNibble)
  {
    int temp = grayNibble & 0x0F;
    temp ^= (temp >> 1);
    temp ^= (temp >> 2);
    temp &= 0x0F;
    return (byte)((grayNibble & 0xF0) | temp);
  }

  private void ProcessTiles()
  {
    // Three reserved tiles:
    //  1. Size. The size of the tiles, in pixels.
    //  2. Length. The length of the payload, in subpixels.
    //  3. Checksum. The checksum of the size, length, and first third of the
    //      payload. The payload is sent in triplicate to allow for some
    //      forward error correction.
    const int kNumReservedTiles = 3;

    // Get the tile size.
    {
      int oldTileSize = tileSize;
      tileSize = kMinTileSize;
      tileSize = Parse24BitTile(0);
      tileSize = Mathf.Clamp(tileSize, kMinTileSize, kMaxTileSize);
      if (tileSize != oldTileSize) {
        Debug.Log($"Tile size changed from {oldTileSize} to {tileSize}");
      }
    }

    // Get the length. This is in units of subpixels. So we will need to access
    // ceil(length/3) tiles.
    int lengthSubpixels = Parse24BitTile(1);
    int lengthTiles = (int) Mathf.Ceil(lengthSubpixels/3.0f);

    // Get the checksum. This covers the size tile, length tile, and first
    // third of the payload.
    int checksumRemote = Parse24BitTile(2);
    int checksumLocal = tileSize + lengthSubpixels;

    Color32 parsed_first = GetTileRGB(0);
    /*
    Debug.Log($"First tile: {parsed_first.r} {parsed_first.g} {parsed_first.b}");
    Debug.Log($"Parsed size {tileSize}");
    Debug.Log($"Parsed length {lengthSubpixels}");
    Debug.Log($"Parsed checksum {checksumRemote}");
    */

    // Collect all nibbles into a flat array. Note that these are still
    // encoded.
    int[] nibbles = new int[lengthSubpixels];
    int nibbleCount = 0;
    for (int tile_i = 0; tile_i < lengthTiles; tile_i++) {
      Color32 parsed_i = GetTileRGB(tile_i+kNumReservedTiles);
      nibbles[nibbleCount++] = parsed_i.r;
      if (nibbleCount < lengthSubpixels) {
        nibbles[nibbleCount++] = parsed_i.g;
      }
      if (nibbleCount < lengthSubpixels) {
        nibbles[nibbleCount++] = parsed_i.b;
      }
    }

    // Compute checksum of nibbles. Match behavior in OperatorEncoding ::
    // Checksum. Note that we only look at the first third of our data,
    // since data is sent in triplicate.
    for (int i = 0; i < nibbleCount / 3; i++) {
      checksumLocal += (nibbles[i] >> 4) & 0x0F;
    }

    //Debug.Log($"Local checksum {checksumLocal}");

    if (checksumLocal != checksumRemote) {
      //Debug.LogWarning($"Checksums don't match. Attempting error recovery.");

      // Data is submitted in triplicate. Perform a bitwise majority vote
      // with `(a & b) | (a & c) | (b & c)`.
      int nc3 = nibbleCount / 3;
      checksumLocal = tileSize + lengthSubpixels;
      for (int i = 0; i < nibbleCount / 3; i++) {
        int copy0 = nibbles[i];
        int copy1 = nibbles[i+nc3];
        int copy2 = nibbles[i+nc3*2];

        // Convert to gray codes before error correction to (hopefully)
        // minimize the number of bit flips which must be corrected.
        byte copy0Gray = ToGrayNibble((byte)((copy0 >> 4) & 0x0F));
        byte copy1Gray = ToGrayNibble((byte)((copy1 >> 4) & 0x0F));
        byte copy2Gray = ToGrayNibble((byte)((copy2 >> 4) & 0x0F));

        int resolvedGray =
          (copy0Gray & copy1Gray) |
          (copy0Gray & copy2Gray) |
          (copy1Gray & copy2Gray);

        byte resolved = ToTwosComplementNibble((byte)resolvedGray);
        nibbles[i] = (resolved << 4) & 0xF0;

        checksumLocal += (nibbles[i] >> 4) & 0x0F;
      }

      // Check result
      if (checksumLocal != checksumRemote) {
        Debug.LogError($"Checksums still don't match after recovery: " +
            $"{checksumRemote} vs {checksumLocal}. Bailing out.");
        return;
      }
    }

    // Convert nibbles to bytes.
    int byteCount = nibbleCount / 6;
    byte[] bytes = new byte[byteCount];
    for (int i = 0; i < byteCount; i++) {
      // See DataEncoder.cs. It puts the upper 4 bits before the lower 4 bits.
      bytes[i] = (byte) ((nibbles[2*i] & 0xF0) | ((nibbles[2*i+1] & 0xF0) >> 4));
    }
    //Debug.Log($"Parsed {bytes.Length} bytes from {nibbles.Length} subpixels");

    // Parse input.
    int bOff = 0;
    while (HasBytesLeft(bytes, bOff, 8)) {
      int type = GetInt(ref bytes, ref bOff);
      int length = GetInt(ref bytes, ref bOff);
      // Can't descend into value if there's not enough length....
      if (!HasBytesLeft(bytes, bOff, length)) {
        break;
      }
      switch (type) {
        case kT_TimeSyncData:
          {
            float syncTimeMs = GetFloat(ref bytes, ref bOff);
            float measureTime = GetFloat(ref bytes, ref bOff) * 1e-6f;
            //Debug.Log($"Parsed time sync data: {syncTimeMs} {measureTimeUs}");

            if (logicalSyncTimeMs != syncTimeMs) {
              // Indicate that we have seen a sync event.
              wallSyncTime = Time.time;
              Debug.Log($"Sync time updated: t0={logicalSyncTimeMs} ms, k=${measureTime}");
            }
            logicalSyncTimeMs = syncTimeMs;
            logicalTimeFactor = 1.0f / measureTime;
            break;
          }
      }
    }
  }

  private int Parse24BitTile(int tileIdx)
  {
    Color32 parsed = GetTileRGB(tileIdx);
    int data = 0;
    data |= DecodeNibble(parsed.r);
    data |= DecodeNibble(parsed.g) << 4;
    data |= DecodeNibble(parsed.b) << 8;
    return data;
  }

  private int DecodeNibble(int subpixel) {
    return (subpixel >> 4) & 0x0F;
  }

  private Color32 GetTileRGB(int tileIdx)
  {
    // Calculate which column and position within column this tile is in
    int tilesPerColumn = readHeight / tileSize;
    int column = tileIdx / Math.Max(1, tilesPerColumn);
    int tileInColumn = tileIdx % tilesPerColumn;

    // Calculate Y position (vertical position within column)
    int tileY = tileInColumn * tileSize;
    int centerY = tileY + tileSize / 2;

    // Calculate X position (horizontal position based on column)
    int tileX = column * tileSize;
    int centerX = tileX + tileSize / 2;

    if (centerY >= readHeight) return new Color32();
    if (centerX >= readWidth) return new Color32();

    int localX = centerX;
    int localY = readHeight - 1 - centerY;
    int index = localY * readWidth + localX;

    if (index < 0 || index >= pixelData.Length) return new Color32();

    return pixelData[index];
  }
}