// SPDX-License-Identifier: AGPL-3.0-or-later
// Copyright (C) 2022-2026 Aykut Alparslan KOC

using System.Diagnostics;
using System.Text;
using Cassandra;
using Ferrite.Data.Repositories;

namespace Ferrite.Data.Repositories;

public class CassandraKVStore : IKVStore
{
    private static readonly TimeSpan DefaultSchemaVisibilityTimeout = TimeSpan.FromSeconds(10);
    private static readonly TimeSpan SchemaVisibilityPollInterval = TimeSpan.FromMilliseconds(50);
    private static long _lastTimestampMicros;
    private readonly ICassandraContext _context;
    private readonly Action<Statement> _enqueue;
    private readonly TimeSpan _schemaVisibilityTimeout;
    private TableDefinition _table;
    private const string IntStr = "int";
    private const string BoolStr = "boolean";
    private const string LongStr = "bigint";
    private const string FloatStr = "float";
    private const string DoubleStr = "double";
    private const string DateStr = "date";
    private const string StringStr = "text";
    private const string BytesStr = "blob";
    private static string GetTypeStr(DataType type) => type switch
    {
        DataType.Bool => BoolStr,
        DataType.Int => IntStr,
        DataType.Long => LongStr,
        DataType.Float => FloatStr,
        DataType.Double => DoubleStr,
        DataType.DateTime => DateStr,
        DataType.String => StringStr,
        DataType.Bytes => BytesStr,
        _ => ""
    };
    private static Type GetManagedType(DataType type) => type switch
    {
        DataType.Bool => typeof(bool),
        DataType.Int => typeof(int),
        DataType.Long => typeof(long),
        DataType.Float => typeof(float),
        DataType.Double => typeof(double),
        DataType.DateTime => typeof(DateTime),
        DataType.String => typeof(string),
        DataType.Bytes => typeof(byte[]),
        _ => typeof(object)
    };
    public CassandraKVStore(ICassandraContext context,
        TimeSpan? schemaVisibilityTimeout = null)
    {
        _context = context;
        _enqueue = context.Enqueue;
        _schemaVisibilityTimeout = schemaVisibilityTimeout ?? DefaultSchemaVisibilityTimeout;
    }

    public CassandraKVStore(ICassandraContext context,
        IWriteBatchAccessor writeBatches,
        TimeSpan? schemaVisibilityTimeout = null)
    {
        _context = context;
        _enqueue = writeBatches.Enqueue;
        _schemaVisibilityTimeout = schemaVisibilityTimeout ?? DefaultSchemaVisibilityTimeout;
    }

    public void SetSchema(TableDefinition table)
    {
        _table = table;
        StringBuilder sb = new StringBuilder($"CREATE TABLE IF NOT EXISTS {_table.Keyspace}.{_table.Name} (");
        bool first = true;
        int pcount = 0;
        foreach (var c in _table.PrimaryKey.Columns)
        {
            pcount++;
            if (!first)
            {
                sb.Append(", ");
            }
            first = false;
            sb.Append($"{c.Name} {GetTypeStr(c.Type)}");
        }

        if (pcount > 0)
        {
            sb.Append(", ");
        }
        pcount = 0;
        sb.Append($"{_table.Name}_data blob, ");
        sb.Append("PRIMARY KEY (");
        first = true;
        foreach (var c in _table.PrimaryKey.Columns)
        {
            if (!first)
            {
                sb.Append(", ");
            }
            first = false;
            sb.Append($"{c.Name}");
        }
        sb.Append("));");
        var statement = new SimpleStatement(sb.ToString());
        CreateTable(_table.Name, statement);
        foreach (var sc in _table.SecondaryIndices)
        {
            pcount = 0;
            sb = new StringBuilder($"CREATE TABLE IF NOT EXISTS {_table.Keyspace}.{_table.Name}_{sc.Name} (");
            first = true;
            foreach (var c in sc.Columns)
            {
                pcount++;
                if (!first)
                {
                    sb.Append(", ");
                }
                first = false;
                sb.Append($"{c.Name} {GetTypeStr(c.Type)}");
            }
            first = true;
            if (pcount > 0)
            {
                sb.Append(", ");
            }
            pcount = 0;
            foreach (var c in _table.PrimaryKey.Columns)
            {
                pcount++;
                if (!first)
                {
                    sb.Append(", ");
                }
                first = false;
                sb.Append($"pk_{c.Name} {GetTypeStr(c.Type)}");
            }
            if (pcount > 0)
            {
                sb.Append(", ");
            }
            sb.Append("PRIMARY KEY (");
            first = true;
            foreach (var c in sc.Columns)
            {
                if (!first)
                {
                    sb.Append(", ");
                }
                first = false;
                sb.Append($"{c.Name}");
            }
            sb.Append("));");
            statement = new SimpleStatement(sb.ToString());
            CreateTable($"{_table.Name}_{sc.Name}", statement);
        }
    }

    private void CreateTable(string name, SimpleStatement statement)
    {
        if (_context.TableIsQueryable(_table.Keyspace, name))
        {
            return;
        }
        ExecuteSchema(statement);
        var waited = Stopwatch.StartNew();
        while (!_context.TableIsQueryable(_table.Keyspace, name))
        {
            if (waited.Elapsed >= _schemaVisibilityTimeout)
            {
                throw new InvalidOperationException(
                    $"{_table.Keyspace}.{name} was created but is still not queryable after " +
                    $"{waited.Elapsed.TotalSeconds:0.###}s; concurrent schema changes have not settled");
            }
            Thread.Sleep(SchemaVisibilityPollInterval);
        }
    }

    private void ExecuteSchema(SimpleStatement statement)
    {
        try
        {
            _context.Execute(statement);
        }
        catch (TimeoutException)
        {
            _context.Execute(new SimpleStatement(statement.QueryString));
        }
    }

    public bool Put(byte[] data, params object[] keys)
    {
        if (keys.Length != _table.PrimaryKey.Columns.Count)
        {
            throw new Exception("Parameter count mismatch.");
        }
        StringBuilder sb = new StringBuilder($"UPDATE {_table.Keyspace}.{_table.Name} " +
                                             $"USING TIMESTAMP ? SET {_table.Name}_data = ? ");
        sb.Append($"WHERE ");
        bool first = true;
        for (int i = 0; i < keys.Length; i++)
        {
            var col = _table.PrimaryKey.Columns[i];
            if (keys[i].GetType() != GetManagedType(col.Type))
            {
                throw new Exception($"Expected type was {GetManagedType(col.Type)} and " +
                                    $"the parameter was of type {keys[i].GetType()}");
            }
            if (!first)
            {
                sb.Append($" AND ");
            }
            first = false;
            sb.Append($"{col.Name} = ?");
        }

        List<object> p = new List<object> { NextTimestampMicros() };
        p.Add(data);
        p.AddRange(keys);
        var statement = new SimpleStatement(sb.ToString(), p.ToArray());
        _enqueue(statement);
        
        foreach (var sc in _table.SecondaryIndices)
        {
            first = true;
            List<object> secondaryParams = new();
            foreach (var c in sc.Columns)
            {
                secondaryParams.Add(keys[_table.PrimaryKey.GetOrdinal(c.Name)]);
            }
            sb = new StringBuilder($"UPDATE {_table.Keyspace}.{_table.Name}_{sc.Name} " +
                                   "USING TIMESTAMP ? SET ");
            foreach (var c in _table.PrimaryKey.Columns)
            {
                if (!first)
                {
                    sb.Append(", ");
                }
                first = false;
                sb.Append($"pk_{c.Name} = ?");
            }
            sb.Append($" WHERE ");
            first = true;
            for (int i = 0; i < secondaryParams.Count; i++)
            {
                var col = sc.Columns[i];
                if (!first)
                {
                    sb.Append($" AND ");
                }
                first = false;
                sb.Append($"{col.Name} = ?");
            }
            List<object> p2 = new List<object> { NextTimestampMicros() };
            p2.AddRange(keys);
            p2.AddRange(secondaryParams);
            var indexStatement = new SimpleStatement(sb.ToString(), p2.ToArray());
            _enqueue(indexStatement);
        }

        return true;
    }

    public bool Delete(params object[] keys)
    {
        EnqueueDelete(keys);
        return true;
    }

    public ValueTask<bool> DeleteAsync(params object[] keys)
    {
        EnqueueDelete(keys);
        return new ValueTask<bool>(true);
    }

    private void EnqueueDelete(object[] keys)
    {
        StringBuilder sb = new StringBuilder($"DELETE FROM {_table.Keyspace}.{_table.Name} " +
                                             "USING TIMESTAMP ? WHERE ");
        bool first = true;
        for (int i = 0; i < keys.Length; i++)
        {
            var col = _table.PrimaryKey.Columns[i];
            if (keys[i].GetType() != GetManagedType(col.Type))
            {
                throw new Exception($"Expected type was {GetManagedType(col.Type)} and " +
                                    $"the parameter was of type {keys[i].GetType()}");
            }
            if (!first)
            {
                sb.Append($" AND ");
            }
            first = false;
            sb.Append($"{col.Name} = ?");
        }
        var parameters = new List<object> { NextTimestampMicros() };
        parameters.AddRange(keys);
        var statement = new SimpleStatement(sb.ToString(), parameters.ToArray());
        EnqueueSecondaryIndexDeletes(keys);
        _enqueue(statement);
    }

    private void EnqueueSecondaryIndexDeletes(object[] keys)
    {
        if (_table.SecondaryIndices.Count == 0)
        {
            return;
        }
        if (keys.Length == _table.PrimaryKey.Columns.Count)
        {
            foreach (var sc in _table.SecondaryIndices)
            {
                List<object> secondaryParams = new();
                foreach (var c in sc.Columns)
                {
                    secondaryParams.Add(keys[_table.PrimaryKey.GetOrdinal(c.Name)]);
                }
                EnqueueIndexRowDelete(sc, secondaryParams);
            }
        }
        else
        {
            foreach (var row in SelectRows(keys))
            {
                foreach (var sc in _table.SecondaryIndices)
                {
                    List<object> secondaryParams = new();
                    foreach (var c in sc.Columns)
                    {
                        secondaryParams.Add(row.GetValue(GetManagedType(c.Type), c.Name));
                    }
                    EnqueueIndexRowDelete(sc, secondaryParams);
                }
            }
        }
    }

    private void EnqueueIndexRowDelete(KeyDefinition sc, List<object> secondaryParams)
    {
        StringBuilder sb = new StringBuilder($"DELETE FROM {_table.Keyspace}.{_table.Name}_{sc.Name} " +
                                             "USING TIMESTAMP ? WHERE ");
        bool first = true;
        for (int i = 0; i < secondaryParams.Count; i++)
        {
            var col = sc.Columns[i];
            if (!first)
            {
                sb.Append($" AND ");
            }
            first = false;
            sb.Append($"{col.Name} = ?");
        }
        var parameters = new List<object> { NextTimestampMicros() };
        parameters.AddRange(secondaryParams);
        var statement = new SimpleStatement(sb.ToString(), parameters.ToArray());
        _enqueue(statement);
    }

    private static long NextTimestampMicros()
    {
        long observed;
        long next;
        do
        {
            observed = Volatile.Read(ref _lastTimestampMicros);
            long now = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds() * 1_000;
            next = Math.Max(now, observed + 1);
        }
        while (Interlocked.CompareExchange(ref _lastTimestampMicros, next,
                   observed) != observed);

        return next;
    }

    private IEnumerable<Row> SelectRows(object[] keys)
    {
        StringBuilder sb = new StringBuilder($"SELECT * FROM {_table.Keyspace}.{_table.Name} WHERE ");
        bool first = true;
        for (int i = 0; i < keys.Length; i++)
        {
            var col = _table.PrimaryKey.Columns[i];
            if (keys[i].GetType() != GetManagedType(col.Type))
            {
                throw new Exception($"Expected type was {GetManagedType(col.Type)} and " +
                                    $"the parameter was of type {keys[i].GetType()}");
            }
            if (!first)
            {
                sb.Append($" AND ");
            }
            first = false;
            sb.Append($"{col.Name} = ?");
        }
        var statement = new SimpleStatement(sb.ToString(), keys.ToArray());
        var results = _context.Execute(statement);
        if (results == null)
        {
            return Enumerable.Empty<Row>();
        }
        return results;
    }

    public bool DeleteBySecondaryIndex(string indexName, params object[] keys)
    {
        var sc = _table.SecondaryIndices.FirstOrDefault(x=>x.Name == indexName);
        if (sc == null)
        {
            return false;
        }
        var statement = new SimpleStatement(BuildIndexSelect(sc, keys), keys.ToArray());
        var results = _context.Execute(statement);
        if (results == null)
        {
            return false;
        }
        var row = results.FirstOrDefault();
        if (row == null)
        {
            return false;
        }
        EnqueueDelete(PrimaryKeyFromIndexRow(row));
        return true;
    }

    public async ValueTask<bool> DeleteBySecondaryIndexAsync(string indexName, params object[] keys)
    {
        var sc = _table.SecondaryIndices.FirstOrDefault(x=>x.Name == indexName);
        if (sc == null)
        {
            return false;
        }
        var statement = new SimpleStatement(BuildIndexSelect(sc, keys), keys.ToArray());
        var results = await _context.ExecuteAsync(statement);
        if (results == null)
        {
            return false;
        }
        var row = results.FirstOrDefault();
        if (row == null)
        {
            return false;
        }
        EnqueueDelete(PrimaryKeyFromIndexRow(row));
        return true;
    }

    private string BuildIndexSelect(KeyDefinition sc, object[] keys)
    {
        StringBuilder sb = new StringBuilder($"SELECT * FROM {_table.Keyspace}.{_table.Name}_{sc.Name} WHERE ");
        bool first = true;
        for (int i = 0; i < keys.Length; i++)
        {
            var col = sc.Columns[i];
            if (keys[i].GetType() != GetManagedType(col.Type))
            {
                throw new Exception($"Expected type was {GetManagedType(col.Type)} and " +
                                    $"the parameter was of type {keys[i].GetType()}");
            }
            if (!first)
            {
                sb.Append($" AND ");
            }
            first = false;
            sb.Append($"{col.Name} = ?");
        }
        return sb.ToString();
    }

    private object[] PrimaryKeyFromIndexRow(Row row)
    {
        List<object> primaryParameters = new();
        foreach (var c in _table.PrimaryKey.Columns)
        {
            primaryParameters.Add(row.GetValue(GetManagedType(c.Type), $"pk_{c.Name}"));
        }
        return primaryParameters.ToArray();
    }

    public byte[]? Get(params object[] keys)
    {
        StringBuilder sb = new StringBuilder($"SELECT * FROM {_table.Keyspace}.{_table.Name} WHERE ");
        bool first = true;
        for (int i = 0; i < keys.Length; i++)
        {
            var col = _table.PrimaryKey.Columns[i];
            if (keys[i].GetType() != GetManagedType(col.Type))
            {
                throw new Exception($"Expected type was {GetManagedType(col.Type)} and " +
                                    $"the parameter was of type {keys[i].GetType()}");
            }
            if (!first)
            {
                sb.Append($" AND ");
            }
            first = false;
            sb.Append($"{col.Name} = ?");
        }
        var statement = new SimpleStatement(sb.ToString(), keys.ToArray());
        var results = _context.Execute(statement);
        if (results == null)
        {
            return null;
        }
        var row = results.FirstOrDefault();
        return row?.GetValue<byte[]>($"{_table.Name}_data");
    }

    public async ValueTask<byte[]?> GetAsync(params object[] keys)
    {
        StringBuilder sb = new StringBuilder($"SELECT * FROM {_table.Keyspace}.{_table.Name} WHERE ");
        bool first = true;
        for (int i = 0; i < keys.Length; i++)
        {
            var col = _table.PrimaryKey.Columns[i];
            if (keys[i].GetType() != GetManagedType(col.Type))
            {
                throw new Exception($"Expected type was {GetManagedType(col.Type)} and " +
                                    $"the parameter was of type {keys[i].GetType()}");
            }
            if (!first)
            {
                sb.Append($" AND ");
            }
            first = false;
            sb.Append($"{col.Name} = ?");
        }
        var statement = new SimpleStatement(sb.ToString(), keys.ToArray());
        var results = await _context.ExecuteAsync(statement);
        if (results == null)
        {
            return null;
        }
        var row = results.FirstOrDefault();
        return row?.GetValue<byte[]>($"{_table.Name}_data");
    }

    public byte[]? GetBySecondaryIndex(string indexName, params object[] keys)
    {
        var sc = _table.SecondaryIndices.FirstOrDefault(x=>x.Name == indexName);
        if (sc != null)
        {
            StringBuilder sb = new StringBuilder($"SELECT * FROM {_table.Keyspace}.{_table.Name}_{sc.Name} WHERE ");
            bool first = true;
            for (int i = 0; i < keys.Length; i++)
            {
                var col = sc.Columns[i];
                if (keys[i].GetType() != GetManagedType(col.Type))
                {
                    throw new Exception($"Expected type was {GetManagedType(col.Type)} and " +
                                        $"the parameter was of type {keys[i].GetType()}");
                }
                if (!first)
                {
                    sb.Append($" AND ");
                }
                first = false;
                sb.Append($"{col.Name} = ?");
            }
            var statement = new SimpleStatement(sb.ToString(), keys.ToArray());
            var results = _context.Execute(statement);
            if (results == null)
            {
                return null;
            }
            var row = results.FirstOrDefault();
            if (row != null)
            {
                var statementInner = GenerateInnerStatement(row);
                var resultsInner = _context.Execute(statementInner);
                var rowInner = resultsInner.FirstOrDefault();
                return rowInner?.GetValue<byte[]>($"{_table.Name}_data");
            }
        }
        return null;
    }

    public async ValueTask<byte[]?> GetBySecondaryIndexAsync(string indexName, params object[] keys)
    {
        var sc = _table.SecondaryIndices.FirstOrDefault(x=>x.Name == indexName);
        if (sc != null)
        {
            StringBuilder sb = new StringBuilder($"SELECT * FROM {_table.Keyspace}.{_table.Name}_{sc.Name} WHERE ");
            bool first = true;
            for (int i = 0; i < keys.Length; i++)
            {
                var col = sc.Columns[i];
                if (keys[i].GetType() != GetManagedType(col.Type))
                {
                    throw new Exception($"Expected type was {GetManagedType(col.Type)} and " +
                                        $"the parameter was of type {keys[i].GetType()}");
                }
                if (!first)
                {
                    sb.Append($" AND ");
                }
                first = false;
                sb.Append($"{col.Name} = ?");
            }
            var statement = new SimpleStatement(sb.ToString(), keys.ToArray());
            var results = await _context.ExecuteAsync(statement);
            if (results == null)
            {
                return null;
            }
            var row = results.FirstOrDefault();
            if (row != null)
            {
                SimpleStatement statementInner = GenerateInnerStatement(row);
                var resultsInner = await _context.ExecuteAsync(statementInner);
                var rowInner = resultsInner.FirstOrDefault();
                return rowInner?.GetValue<byte[]>($"{_table.Name}_data");
            }
        }
        return null;
    }

    private SimpleStatement GenerateInnerStatement(Row row)
    {
        StringBuilder sb;
        bool first;
        List<object> primaryParameters = new();
        foreach (var c in _table.PrimaryKey.Columns)
        {
            primaryParameters.Add(row.GetValue(GetManagedType(c.Type), $"pk_{c.Name}"));
        }

        sb = new StringBuilder($"SELECT * FROM {_table.Keyspace}.{_table.Name} WHERE ");
        first = true;
        for (int i = 0; i < primaryParameters.Count; i++)
        {
            var col = _table.PrimaryKey.Columns[i];
            if (!first)
            {
                sb.Append($" AND ");
            }

            first = false;
            sb.Append($"{col.Name} = ?");
        }

        var statementInner = new SimpleStatement(sb.ToString(), primaryParameters.ToArray());
        return statementInner;
    }

    public IEnumerable<byte[]> Iterate(params object[] keys)
    {
        StringBuilder sb = new StringBuilder(
            $"SELECT * FROM {_table.Keyspace}.{_table.Name}");
        if (keys.Length > 0)
        {
            sb.Append(" WHERE ");
        }
        bool first = true;
        for (int i = 0; i < keys.Length; i++)
        {
            var col = _table.PrimaryKey.Columns[i];
            if (keys[i].GetType() != GetManagedType(col.Type))
            {
                throw new Exception($"Expected type was {GetManagedType(col.Type)} and " +
                                    $"the parameter was of type {keys[i].GetType()}");
            }
            if (!first)
            {
                sb.Append($" AND ");
            }
            first = false;
            sb.Append($"{col.Name} = ?");
        }
        var statement = new SimpleStatement(sb.ToString(), keys.ToArray());
        var results = _context.Execute(statement);
        if (results == null)
        {
            yield break;
        }
        foreach (var row in results)
        {
            yield return row.GetValue<byte[]>($"{_table.Name}_data");
        }
    }
    public async IAsyncEnumerable<byte[]> IterateAsync(params object[] keys)
    {
        StringBuilder sb = new StringBuilder(
            $"SELECT * FROM {_table.Keyspace}.{_table.Name}");
        if (keys.Length > 0)
        {
            sb.Append(" WHERE ");
        }
        bool first = true;
        for (int i = 0; i < keys.Length; i++)
        {
            var col = _table.PrimaryKey.Columns[i];
            if (keys[i].GetType() != GetManagedType(col.Type))
            {
                throw new Exception($"Expected type was {GetManagedType(col.Type)} and " +
                                    $"the parameter was of type {keys[i].GetType()}");
            }
            if (!first)
            {
                sb.Append($" AND ");
            }
            first = false;
            sb.Append($"{col.Name} = ?");
        }
        var statement = new SimpleStatement(sb.ToString(), keys.ToArray());
        var results = await _context.ExecuteAsync(statement);
        if (results == null)
        {
            yield break;
        }
        foreach (var row in results)
        {
            yield return row.GetValue<byte[]>($"{_table.Name}_data");
        }
    }

    public IEnumerable<byte[]> IterateBySecondaryIndex(string indexName, params object[] keys)
    {
        var sc = _table.SecondaryIndices.FirstOrDefault(x=>x.Name == indexName);
        if (sc != null)
        {
            StringBuilder sb = new StringBuilder($"SELECT * FROM {_table.Keyspace}.{_table.Name}_{sc.Name} WHERE ");
            bool first = true;
            for (int i = 0; i < keys.Length; i++)
            {
                var col = sc.Columns[i];
                if (keys[i].GetType() != GetManagedType(col.Type))
                {
                    throw new Exception($"Expected type was {GetManagedType(col.Type)} and " +
                                        $"the parameter was of type {keys[i].GetType()}");
                }
                if (!first)
                {
                    sb.Append($" AND ");
                }
                first = false;
                sb.Append($"{col.Name} = ?");
            }
            var statement = new SimpleStatement(sb.ToString(), keys.ToArray());
            var results = _context.Execute(statement);
            if (results == null)
            {
                yield break;
            }
            foreach (var row in results)
            {
                SimpleStatement statementInner = GenerateInnerStatement(row);
                var resultsInner = _context.Execute(statementInner);
                var rowInner = resultsInner.FirstOrDefault();
                yield return rowInner?.GetValue<byte[]>($"{_table.Name}_data");
            }
        }
    }

    public async IAsyncEnumerable<byte[]> IterateBySecondaryIndexAsync(string indexName, params object[] keys)
    {
        var sc = _table.SecondaryIndices.FirstOrDefault(x=>x.Name == indexName);
        if (sc != null)
        {
            StringBuilder sb = new StringBuilder($"SELECT * FROM {_table.Keyspace}.{_table.Name}_{sc.Name} WHERE ");
            bool first = true;
            for (int i = 0; i < keys.Length; i++)
            {
                var col = sc.Columns[i];
                if (keys[i].GetType() != GetManagedType(col.Type))
                {
                    throw new Exception($"Expected type was {GetManagedType(col.Type)} and " +
                                        $"the parameter was of type {keys[i].GetType()}");
                }
                if (!first)
                {
                    sb.Append($" AND ");
                }
                first = false;
                sb.Append($"{col.Name} = ?");
            }
            var statement = new SimpleStatement(sb.ToString(), keys.ToArray());
            var results = await _context.ExecuteAsync(statement);
            if (results == null)
            {
                yield break;
            }

            foreach (var row in results)
            {
                SimpleStatement statementInner = GenerateInnerStatement(row);
                var resultsInner = await _context.ExecuteAsync(statementInner);
                var rowInner = resultsInner.FirstOrDefault();
                yield return rowInner?.GetValue<byte[]>($"{_table.Name}_data");
            }
        }
    }
}