using System.Text;

namespace NathanMcRae;

/// <summary>
/// Sane Tab-Separated Values
/// </summary>
public class SaneTsv
{
  // Like an enum, but more extensible
  public class ColumnType { }
  public class StringType : ColumnType { }
  public class BooleanType : ColumnType { }
  public class Float32Type : ColumnType { }
  public class Float32LEType : ColumnType { }
  public class Float64Type : ColumnType { }
  public class Float64LEType : ColumnType { }
  public class UInt32Type : ColumnType { }
  public class UInt64Type : ColumnType { }
  public class Int32Type : ColumnType { }
  public class Int64Type : ColumnType { }
  public class BinaryType : ColumnType { }

  protected enum FormatType
  {
    SIMPLE_TSV = 0,
    TYPED_TSV = 1,
    COMMENTED_TSV = 2,
  }

  public static readonly byte[] TrueEncoded = Encoding.UTF8.GetBytes("TRUE");
  public static readonly byte[] FalseEncoded = Encoding.UTF8.GetBytes("FALSE");

  // TODO: We need to be able to update all these in tandem somehow
  public string[] ColumnNames { get; protected set; }
  public Type[] ColumnTypes { get; protected set; }
  public List<SaneTsvRecord> Records { get; protected set; }
  public string FileComment { get; protected set; } = null;
  protected static bool? _littleEndian = null;
  public static bool LittleEndian
  {
    get
    {
      if (_littleEndian == null)
      {
        _littleEndian = BitConverter.GetBytes(double.NegativeInfinity)[7] == 255;
      }
      return _littleEndian.Value;
    }
  }

  public static SaneTsv ParseSimpleTsv(byte[] inputBuffer)
  {
    return Parse(inputBuffer, FormatType.SIMPLE_TSV);
  }

  public static SaneTsv ParseTypedTsv(byte[] inputBuffer)
  {
    return Parse(inputBuffer, FormatType.TYPED_TSV);
  }

  public static SaneTsv ParseCommentedTsv(byte[] inputBuffer)
  {
    return Parse(inputBuffer, FormatType.COMMENTED_TSV);
  }

  // TODO: Have parsing errors include line / column #
  protected static SaneTsv Parse(byte[] inputBuffer, FormatType format)
  {
    var parsed = new SaneTsv();
    parsed.ColumnNames = new string[] { };
    parsed.ColumnTypes = new Type[] { };
    parsed.Records = new List<SaneTsvRecord>();

    var fieldBytes = new List<byte>();
    var fields = new List<byte[]>();
    var currentComment = new StringBuilder();

    int numFields = -1;
    int line = 1;
    int currentLineStart = 0;
    for (int i = 0; i < inputBuffer.Count(); i++)
    {
      if (inputBuffer[i] == '\\')
      {
        if (i + 1 == inputBuffer.Count())
        {
          throw new Exception($"Found '\\' at end of input");
        }
        if (inputBuffer[i + 1] == 'n')
        {
          fieldBytes.Add((byte)'\n');
          i++;
        }
        else if (inputBuffer[i + 1] == '\\')
        {
          fieldBytes.Add((byte)'\\');
          i++;
        }
        else if (inputBuffer[i + 1] == 't')
        {
          fieldBytes.Add((byte)'\t');
          i++;
        }
        else if (inputBuffer[i + 1] == '#')
        {
          fieldBytes.Add((byte)'#');
          i++;
        }
        else
        {
          throw new Exception($"Expected 'n', 't', '#', or '\\' after '\\' at line {line} column {i - currentLineStart}");
        }
      }
      else if (inputBuffer[i] == '\t')
      {
        // end of field
        fields.Add(fieldBytes.ToArray());
        fieldBytes.Clear();
      }
      else if (inputBuffer[i] == '\n')
      {
        fields.Add(fieldBytes.ToArray());
        fieldBytes.Clear();

        if (numFields < 0)
        {
          // This is the header

          numFields = fields.Count;

          parsed.ColumnNames = new string[numFields];
          parsed.ColumnTypes = new Type[numFields];

          int numTypesBlank = 0;

          for (int j = 0; j < fields.Count; j++)
          {
            string columnString;
            try
            {
              columnString = Encoding.UTF8.GetString(fields[j]);
            }
            catch (Exception e)
            {
              throw new Exception($"Header {fields.Count} is not valid UTF-8", e);
            }

            string columnTypeString;
            string columnName;
            if (columnString.Contains(':'))
            {
              if (format == FormatType.SIMPLE_TSV)
              {
                throw new Exception($"Header {fields.Count} contain ':', which is not allowed for column names");
              }
              columnTypeString = columnString.Split(":").Last();
              columnName = columnString.Substring(0, columnString.Length - columnTypeString.Length - 1);
            }
            else
            {
              if (format > FormatType.SIMPLE_TSV)
              {
                throw new Exception($"Header {fields.Count} has no type");
              }
              columnTypeString = "";
              columnName = columnString;
            }

            Type type;

            switch (columnTypeString)
            {
              case "":
                numTypesBlank++;
                type = typeof(StringType);
                break;
              case "string":
                type = typeof(StringType);
                break;
              case "boolean":
                type = typeof(BooleanType);
                break;
              case "float32":
                type = typeof(Float32Type);
                break;
              case "float32-le":
                type = typeof(Float32LEType);
                break;
              case "float64":
                type = typeof(Float64Type);
                break;
              case "float64-le":
                type = typeof(Float64LEType);
                break;
              case "uint32":
                type = typeof(UInt32Type);
                break;
              case "uint64":
                type = typeof(UInt64Type);
                break;
              case "int32":
                type = typeof(Int32Type);
                break;
              case "int64":
                type = typeof(Int64Type);
                break;
              case "binary":
                type = typeof(BinaryType);
                break;
              default:
                throw new Exception($"Invalid type '{columnTypeString}' for column {j}");
            }

            // TODO: Check column name uniqueness

            parsed.ColumnNames[j] = columnName;
            parsed.ColumnTypes[j] = type;
          }

          if (currentComment.Length > 0)
          {
            parsed.FileComment = currentComment.ToString();
            currentComment.Clear();
          }

          fields.Clear();
        }
        else if (numFields != fields.Count)
        {
          throw new Exception($"Expected {numFields} fields on line {line}, but found {fields.Count}");
        }
        else
        {
          string comment = null;
          if (currentComment.Length > 0)
          {
            comment = currentComment.ToString();
            currentComment.Clear();
          }
          parsed.Records.Add(new SaneTsvRecord(parsed.ColumnNames, ParseCurrentRecord(parsed, fields, line), comment, line));
          fields.Clear();
        }

        line++;
        currentLineStart = i + 1;
      }
      else if (inputBuffer[i] == '#')
      {
        if (i == currentLineStart && format >= FormatType.COMMENTED_TSV)
        {
          int j = i;
          for (; j < inputBuffer.Length && inputBuffer[j] != '\n'; j++) { }
          if (j < inputBuffer.Length)
          {
            var commentBytes = new byte[j - i - 1];
            Array.Copy(inputBuffer, i + 1, commentBytes, 0, j - i - 1);
            if (currentComment.Length > 0)
            {
              currentComment.Append('\n');
            }
            currentComment.Append(Encoding.UTF8.GetString(commentBytes));
            i = j;
            currentLineStart = i + 1;
            line++;
          }
          else
          {
            throw new Exception("Comments at end of file are not allowed");
          }
        }
        else
        {
          throw new Exception($"Found unescaped '#' at line {line}, column {i - currentLineStart}");
        }
      }
      else
      {
        fieldBytes.Add(inputBuffer[i]);
      }
    }

    fields.Add(fieldBytes.ToArray());

    if (numFields == 0)
    {
      throw new Exception("Found 0 fields on last line. Possibly because of extra \\n after last record");
    }
    if (numFields != fields.Count)
    {
      throw new Exception($"Expected {numFields} fields on line {parsed.Records.Count + 2}, but found {fields.Count}");
    }
    else
    {
      string comment = null;
      if (currentComment.Length > 0)
      {
        comment = currentComment.ToString();
        currentComment.Clear();
      }
      parsed.Records.Add(new SaneTsvRecord(parsed.ColumnNames, ParseCurrentRecord(parsed, fields, line), comment, line));
      fields.Clear();
    }

    return parsed;
  }

  /// <summary>
  /// Note: this modifies 'parsed'
  /// </summary>
  protected static object[] ParseCurrentRecord(SaneTsv parsed, List<byte[]> fields, int line)
  {
    var parsedFields = new object[fields.Count];
    for (int j = 0; j < fields.Count; j++)
    {
      // All other types require the content to be UTF-8. Binary fields can ignore that.
      if (parsed.ColumnTypes[j] == typeof(BinaryType))
      {
        parsedFields[j] = fields[j];
        continue;
      }
      else if (parsed.ColumnTypes[j] == typeof(Float32LEType))
      {
        byte[] floatBytes;
        if (!LittleEndian)
        {
          floatBytes = new byte[sizeof(float)];
          for (int k = 0; k < sizeof(float); k++)
          {
            floatBytes[k] = fields[j][sizeof(float) - 1 - k];
          }
        }
        else
        {
          floatBytes = fields[j];
        }
        parsedFields[j] = BitConverter.ToSingle(floatBytes, 0);

        continue;
      }
      else if (parsed.ColumnTypes[j] == typeof(Float64LEType))
      {
        byte[] floatBytes;
        if (!LittleEndian)
        {
          floatBytes = new byte[sizeof(double)];
          for (int k = 0; k < sizeof(double); k++)
          {
            floatBytes[k] = fields[j][sizeof(double) - 1 - k];
          }
        }
        else
        {
          floatBytes = fields[j];
        }
        parsedFields[j] = BitConverter.ToDouble(floatBytes, 0);

        continue;
      }

      string fieldString;
      try
      {
        fieldString = Encoding.UTF8.GetString(fields[j]);
      }
      catch (Exception e)
      {
        throw new Exception($"Field {j} on line {line} is not valid UTF-8", e);
      }

      // TODO: Add checking for numeric types format

      if (parsed.ColumnTypes[j] == typeof(StringType))
      {
        parsedFields[j] = fieldString;
      }
      else if (parsed.ColumnTypes[j] == typeof(BooleanType))
      {
        bool parsedBool;
        if (fieldString == "TRUE")
        {
          parsedBool = true;
        }
        else if (fieldString == "FALSE")
        {
          parsedBool = false;
        }
        else
        {
          throw new Exception($"Field {j} on line {line} is not valid boolean. Must be 'TRUE' or 'FALSE' exactly");
        }

        parsedFields[j] = parsedBool;
      }
      else if (parsed.ColumnTypes[j] == typeof(Float32Type))
      {
        float parsedFloat;
        if (!float.TryParse(fieldString, out parsedFloat))
        {
          if (fieldString == "-inf")
          {
            parsedFloat = float.NegativeInfinity;
          }
          else if (fieldString == "+inf")
          {
            parsedFloat = float.PositiveInfinity;
          }
          else
          {
            throw new Exception($"Field {j} on line {line} is not valid single-precision float");
          }
        }

        parsedFields[j] = parsedFloat;
      }
      else if (parsed.ColumnTypes[j] == typeof(Float64Type))
      {
        double parsedDouble;
        if (!double.TryParse(fieldString, out parsedDouble))
        {
          if (fieldString == "-inf")
          {
            parsedDouble = float.NegativeInfinity;
          }
          else if (fieldString == "+inf")
          {
            parsedDouble = float.PositiveInfinity;
          }
          else
          {
            throw new Exception($"Field {j} on line {line} is not valid double-precision float");
          }
        }

        parsedFields[j] = parsedDouble;
      }
      else if (parsed.ColumnTypes[j] == typeof(UInt32Type))
      {
        if (!UInt32.TryParse(fieldString, out UInt32 parsedUInt32))
        {
          throw new Exception($"Field {j} on line {line} is not valid UInt32");
        }

        parsedFields[j] = parsedUInt32;
      }
      else if (parsed.ColumnTypes[j] == typeof(UInt64Type))
      {
        if (!UInt64.TryParse(fieldString, out UInt64 parsedUInt64))
        {
          throw new Exception($"Field {j} on line {line} is not valid UInt64");
        }

        parsedFields[j] = parsedUInt64;
      }
      else if (parsed.ColumnTypes[j] == typeof(Int32Type))
      {
        if (!Int32.TryParse(fieldString, out Int32 parsedInt32))
        {
          throw new Exception($"Field {j} on line {line} is not valid Int32");
        }

        parsedFields[j] = parsedInt32;
      }
      else if (parsed.ColumnTypes[j] == typeof(Int64Type))
      {
        if (!Int64.TryParse(fieldString, out Int64 parsedInt64))
        {
          throw new Exception($"Field {j} on line {line} is not valid Int64");
        }

        parsedFields[j] = parsedInt64;
      }
      else
      {
        throw new Exception($"Unexpected type {parsed.ColumnTypes[j]}");
      }
    }

    return parsedFields;
  }

  public static byte[] SerializeSimpleTsv(IList<string> header, IList<IList<string>> data)
  {
    var escapedString = new StringBuilder();

    // Serialize header
    for (int i = 0; i < header.Count; i++)
    {
      if (header[i].Contains(':'))
      {
        throw new Exception($"Column {i} contains the character ':'");
      }

      for (int j = i + 1; j < header.Count; j++)
      {
        if (header[i] == header[j])
        {
          throw new Exception("Column names in header must be unique");
        }
      }

      for (int j = 0; j < header[i].Count(); j++)
      {
        if (header[i][j] == '\n')
        {
          escapedString.Append("\\n");
        }
        else if (header[i][j] == '\t')
        {
          escapedString.Append("\\t");
        }
        else if (header[i][j] == '\\')
        {
          escapedString.Append("\\\\");
        }
        else if (header[i][j] == '#')
        {
          escapedString.Append("\\#");
        }
        else
        {
          escapedString.Append(header[i][j]);
        }
      }

      if (i == header.Count - 1)
      {
        escapedString.Append('\n');
      }
      else
      {
        escapedString.Append('\t');
      }
    }

    // Serialize data
    for (int i = 0; i < data.Count; i++)
    {
      for (int j = 0; j < data[i].Count; j++)
      {
        for (int k = 0; k < data[i][j].Length; k++)
        {
          if (data[i][j][k] == '\n')
          {
            escapedString.Append("\\n");
          }
          else if (data[i][j][k] == '\t')
          {
            escapedString.Append("\\t");
          }
          else if (data[i][j][k] == '\\')
          {
            escapedString.Append("\\\\");
          }
          else if (data[i][j][k] == '#')
          {
            escapedString.Append("\\#");
          }
          else
          {
            escapedString.Append(data[i][j][k]);
          }
        }

        if (j < data[i].Count - 1)
        {
          escapedString.Append('\t');
        }
        else if (i < data.Count - 1)
        {
          escapedString.Append('\n');
        }
      }
    }

    return Encoding.UTF8.GetBytes(escapedString.ToString());
  }

  public static Type GetColumnFromType(Type type)
  {
    if (type == typeof(string))
    {
      return typeof(StringType);
    }
    else if (type == typeof(bool))
    {
      return typeof(BooleanType);
    }
    else if (type == typeof(float))
    {
      return typeof(Float32Type);
    }
    else if (type == typeof(double))
    {
      return typeof(Float64Type);
    }
    else if (type == typeof(UInt32))
    {
      return typeof(UInt32Type);
    }
    else if (type == typeof(UInt64))
    {
      return typeof(UInt64Type);
    }
    else if (type == typeof(Int32))
    {
      return typeof(Int32Type);
    }
    else if (type == typeof(Int64))
    {
      return typeof(Int64Type);
    }
    else if (type == typeof(byte[]))
    {
      return typeof(BinaryType);
    }
    else
    {
      throw new Exception($"Invalid type: {type}");
    }
  }

  public static string GetNameFromColumn(Type type)
  {
    if (type == typeof(StringType))
    {
      return "string";
    }
    else if (type == typeof(BooleanType))
    {
      return "boolean";
    }
    else if (type == typeof(Float32Type))
    {
      return "float32";
    }
    else if (type == typeof(Float32LEType))
    {
      return "float32-le";
    }
    else if (type == typeof(Float64Type))
    {
      return "float64";
    }
    else if (type == typeof(Float64LEType))
    {
      return "float64-le";
    }
    else if (type == typeof(UInt32Type))
    {
      return "uint32";
    }
    else if (type == typeof(UInt64Type))
    {
      return "uint64";
    }
    else if (type == typeof(Int32Type))
    {
      return "int32";
    }
    else if (type == typeof(Int64Type))
    {
      return "int64";
    }
    else if (type == typeof(BinaryType))
    {
      return "binary";
    }
    else
    {
      throw new Exception($"Invalid type: {type}");
    }
  }

  public static byte[] SerializeTypedTsv(IList<Type> headerTypes, IList<string> headerNames, IList<IList<object>> data)
  {
    var bytes = new List<byte>();

    if (headerNames.Count != headerTypes.Count)
    {
      throw new ArgumentException($"headerTypes length ({headerTypes.Count}) does not match headerNames length ({headerNames.Count})");
    }

    // Serialize header
    for (int i = 0; i < headerNames.Count; i++)
    {
      for (int j = i + 1; j < headerNames.Count; j++)
      {
        if (headerNames[i] == headerNames[j])
        {
          throw new Exception("Column names in header must be unique");
        }
      }

      byte[] nameEncoded = Encoding.UTF8.GetBytes(headerNames[i]);

      for (int j = 0; j < nameEncoded.Length; j++)
      {
        if (nameEncoded[j] == '\n')
        {
          bytes.Add((byte)'\\');
          bytes.Add((byte)'n');
        }
        else if (nameEncoded[j] == '\t')
        {
          bytes.Add((byte)'\\');
          bytes.Add((byte)'t');
        }
        else if (nameEncoded[j] == '\\')
        {
          bytes.Add((byte)'\\');
          bytes.Add((byte)'\\');
        }
        else if (nameEncoded[j] == '#')
        {
          bytes.Add((byte)'\\');
          bytes.Add((byte)'#');
        }
        else
        {
          bytes.Add(nameEncoded[j]);
        }
      }

      bytes.Add((byte)':');
      try
      {
        bytes.AddRange(Encoding.UTF8.GetBytes(GetNameFromColumn(headerTypes[i])));
      }
      catch (Exception e)
      {
        throw new Exception($"Invalid header type for column {i}", e);
      }

      if (i == headerNames.Count - 1)
      {
        bytes.Add((byte)'\n');
      }
      else
      {
        bytes.Add((byte)'\t');
      }
    }

    // Serialize data
    for (int i = 0; i < data.Count; i++)
    {
      for (int j = 0; j < data[i].Count; j++)
      {
        try
        {
          byte[] fieldEncoded = null;
          // Some fields definitely don't need escaping, so we add them directly to bytes
          bool skipEscaping = false;

          if (headerTypes[j] == typeof(StringType))
          {
            fieldEncoded = Encoding.UTF8.GetBytes((string)data[i][j]);
          }
          else if (headerTypes[j] == typeof(BooleanType))
          {
            bytes.AddRange((bool)data[i][j] ? TrueEncoded : FalseEncoded);
            skipEscaping = true;
          }
          else if (headerTypes[j] == typeof(Float32Type))
          {
            if (data[i][j] is float f)
            {
              if (float.IsNegativeInfinity(f))
              {
                bytes.AddRange(Encoding.UTF8.GetBytes("-inf"));
              }
              else if (float.IsPositiveInfinity(f))
              {
                bytes.AddRange(Encoding.UTF8.GetBytes("+inf"));
              }
              else
              {
                // See https://learn.microsoft.com/en-us/dotnet/standard/base-types/standard-numeric-format-strings#round-trip-format-specifier-r
                bytes.AddRange(Encoding.UTF8.GetBytes(((float)data[i][j]).ToString("G9")));
              }
            }
            else
            {
              throw new InvalidCastException();
            }
            skipEscaping = true;
          }
          else if (headerTypes[j] == typeof(Float32LEType))
          {
            if (LittleEndian)
            {
              fieldEncoded = BitConverter.GetBytes((float)data[i][j]);
            }
            else
            {
              byte[] floatBytes = BitConverter.GetBytes((float)data[i][j]);
              fieldEncoded = new byte[sizeof(float)];
              for (int k = 0; k < sizeof(float); k++)
              {
                fieldEncoded[k] = floatBytes[sizeof(float) - 1 - k];
              }
            }
          }
          else if (headerTypes[j] == typeof(Float64Type))
          {
            if (data[i][j] is double d)
            {
              if (double.IsNegativeInfinity(d))
              {
                bytes.AddRange(Encoding.UTF8.GetBytes("-inf"));
              }
              else if (double.IsPositiveInfinity(d))
              {
                bytes.AddRange(Encoding.UTF8.GetBytes("+inf"));
              }
              else
              {
                // See https://learn.microsoft.com/en-us/dotnet/standard/base-types/standard-numeric-format-strings#round-trip-format-specifier-r
                bytes.AddRange(Encoding.UTF8.GetBytes((d).ToString("G17")));
              }
            }
            else
            {
              throw new InvalidCastException();
            }
            skipEscaping = true;
          }
          else if (headerTypes[j] == typeof(Float64LEType))
          {
            if (LittleEndian)
            {
              fieldEncoded = BitConverter.GetBytes((double)data[i][j]);
            }
            else
            {
              byte[] doubleBytes = BitConverter.GetBytes((double)data[i][j]);
              fieldEncoded = new byte[sizeof(double)];
              for (int k = 0; k < sizeof(double); k++)
              {
                fieldEncoded[k] = doubleBytes[sizeof(double) - 1 - k];
              }
            }
          }
          else if (headerTypes[j] == typeof(UInt32Type))
          {
            bytes.AddRange(Encoding.UTF8.GetBytes(((UInt32)(int)data[i][j]).ToString()));
            skipEscaping = true;
          }
          else if (headerTypes[j] == typeof(UInt64Type))
          {
            bytes.AddRange(Encoding.UTF8.GetBytes(((UInt64)(int)data[i][j]).ToString()));
            skipEscaping = true;
          }
          else if (headerTypes[j] == typeof(Int32Type))
          {
            bytes.AddRange(Encoding.UTF8.GetBytes(((Int32)(int)data[i][j]).ToString()));
            skipEscaping = true;
          }
          else if (headerTypes[j] == typeof(Int64Type))
          {
            bytes.AddRange(Encoding.UTF8.GetBytes(((Int64)(int)data[i][j]).ToString()));
            skipEscaping = true;
          }
          else if (headerTypes[j] == typeof(BinaryType))
          {
            fieldEncoded = (byte[])data[i][j];
          }
          else
          {
            throw new Exception($"Unexpected column type {headerTypes[j]} for column {j}");
          }

          if (!skipEscaping)
          {
            for (int k = 0; k < fieldEncoded.Length; k++)
            {
              if (fieldEncoded[k] == '\n')
              {
                bytes.Add((byte)'\\');
                bytes.Add((byte)'n');
              }
              else if (fieldEncoded[k] == '\t')
              {
                bytes.Add((byte)'\\');
                bytes.Add((byte)'t');
              }
              else if (fieldEncoded[k] == '\\')
              {
                bytes.Add((byte)'\\');
                bytes.Add((byte)'\\');
              }
              else if (fieldEncoded[k] == '#')
              {
                bytes.Add((byte)'\\');
                bytes.Add((byte)'#');
              }
              else
              {
                bytes.Add(fieldEncoded[k]);
              }
            }
          }

          if (j < data[i].Count - 1)
          {
            bytes.Add((byte)'\t');
          }
          else if (i < data.Count - 1)
          {
            bytes.Add((byte)'\n');
          }
        }
        catch (InvalidCastException e)
        {
          throw new Exception($"Record {i}, field {j} expected type compatible with {GetNameFromColumn(headerTypes[j])}", e);
        }
      }
    }

    return bytes.ToArray();
  }

  public SaneTsvRecord this[int i] => Records[i];

  public class SaneTsvRecord
  {
    public string[] ColumnNames { get; }
    public string Comment { get; }
    public object[] Fields { get; }
    public int Line { get; }

    public object this[string columnName] => Fields[Array.IndexOf(ColumnNames, columnName)];

    public object this[int columnIndex] => Fields[columnIndex];

    public SaneTsvRecord(string[] columnNames, object[] fields, string comment, int line)
    {
      ColumnNames = columnNames;
      Fields = fields;
      Comment = comment;
      Line = line;
    }
  }
}