// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

namespace TypeAgent.KnowPro.Storage.Sqlite;

// TODO: update methods to use new wrappers and extension methods
//  in SqliteDatabase etc. Should not need to use ExecuteReader
// directly in most cases

public class SqliteSemanticRefCollection : ISemanticRefCollection
{
    SqliteDatabase _db;
    int _count = -1;

    public SqliteSemanticRefCollection(SqliteDatabase db)
    {
        ArgumentVerify.ThrowIfNull(db, nameof(db));
        _db = db;
    }


    public bool IsPersistent => true;

    public int GetCount()
    {
        if (_count < 0)
        {
            _count = _db.GetCount(SqliteStorageProviderSchema.SemanticRefTable);
        }
        return _count;
    }

    public ValueTask<int> GetCountAsync(CancellationToken cancellationToken = default)
    {
        return ValueTask.FromResult(GetCount());
    }

    public void Append(SemanticRef semanticRef)
    {
        KnowProVerify.ThrowIfInvalid(semanticRef);

        SemanticRefRow row = ToSemanticRefRow(semanticRef);

        using var cmd = _db.CreateCommand(
           @"INSERT INTO SemanticRefs (semref_id, range_json, knowledge_type, knowledge_json)
          VALUES (@semref_id, @range_json, @knowledge_type, @knowledge_json);"
        );
        WriteSemanticRefRow(cmd, row);
        int rowCount = cmd.ExecuteNonQuery();
        if (rowCount > 0)
        {
            _count += rowCount;
        }
    }

    public ValueTask AppendAsync(SemanticRef semanticRef, CancellationToken cancellationToken = default)
    {
        Append(semanticRef);
        return ValueTask.CompletedTask;
    }

    public ValueTask AppendAsync(IList<SemanticRef> items, CancellationToken cancellationToken = default)
    {
        ArgumentVerify.ThrowIfNull(items, nameof(items));

        // TODO: Bulk operations
        foreach (var sr in items)
        {
            cancellationToken.ThrowIfCancellationRequested();
            Append(sr);
        }
        return ValueTask.CompletedTask;
    }

    public SemanticRef Get(int semanticRefId)
    {
        ArgumentVerify.ThrowIfLessThan(semanticRefId, 0, nameof(semanticRefId));

        using var cmd = _db.CreateCommand(@"
SELECT semref_id, range_json, knowledge_type, knowledge_json
FROM SemanticRefs WHERE semref_id = @semref_id");
        cmd.AddParameter("semref_id", semanticRefId);

        using var reader = cmd.ExecuteReader();
        if (!reader.Read())
        {
            throw new ArgumentException($"No semanticRef at ordinal {semanticRefId}");
        }

        SemanticRefRow row = ReadSemanticRefRow(reader);
        SemanticRef semanticRef = FromSemanticRefRow(row);
        return semanticRef;
    }


    public ValueTask<SemanticRef> GetAsync(int semanticRefId, CancellationToken cancellationToken = default)
    {
        return ValueTask.FromResult(Get(semanticRefId));
    }

    public ValueTask<IList<SemanticRef>> GetAsync(IList<int> semanticRefIds, CancellationToken cancellationToken = default)
    {
        ArgumentVerify.ThrowIfNullOrEmpty(semanticRefIds, nameof(semanticRefIds));

        Dictionary<int, SemanticRef> dicSemanticRefs = [];
        foreach (var batch in semanticRefIds.Batch(SqliteDatabase.MaxBatchSize))
        {
            // Since the database doesn't gurantee the order of returned rows within IN statements
            // rows come back based on the index used by the query or by order the were inserted)
            // or some other unguarnateed order
            // We use the dictionary to retain the order and to do easily lookups from the db results
            foreach (var item in batch)
            {
                dicSemanticRefs.Add(item, SemanticRef.Empty);
            }

            var placeholderIds = SqliteDatabase.MakeInPlaceholderParamIds(batch.Count);
            var sql = $@"
SELECT semref_id, range_json, knowledge_type, knowledge_json
FROM SemanticRefs WHERE semref_id IN ({SqliteDatabase.MakeInStatement(placeholderIds)})
ORDER BY semref_id";
            var rows = _db.Enumerate(
                sql,
                cmd => cmd.AddPlaceholderParameters(placeholderIds, batch),
                ReadSemanticRefRow
            );
            foreach (var row in rows)
            {
                var sr = FromSemanticRefRow(row);
                dicSemanticRefs[sr.SemanticRefOrdinal] = sr;
            }
        }
        return ValueTask.FromResult((IList<SemanticRef>)[.. dicSemanticRefs.Values]);
    }

    public TextRange GetRange(int semanticRefId)
    {
        ArgumentVerify.ThrowIfLessThan(semanticRefId, 0, nameof(semanticRefId));

        return _db.Get(@"
SELECT range_json
FROM SemanticRefs WHERE semref_id = @semref_id",
            (cmd) => cmd.AddParameter("semref_id", semanticRefId),
            (reader) => StorageSerializer.FromJson<TextRange>(reader.GetString(0))
        );
    }

    public ValueTask<TextRange> GetTextRangeAsync(int semanticRefId, CancellationToken cancellationToken = default)
    {
        return ValueTask.FromResult(GetRange(semanticRefId));
    }

    public ValueTask<IList<TextRange>> GetTextRangeAsync(IList<int> semanticRefIds, CancellationToken cancellationToken = default)
    {
        ArgumentVerify.ThrowIfNullOrEmpty(semanticRefIds, nameof(semanticRefIds));

        Dictionary<int, TextRange> rangeRows = [];
        foreach (var batch in semanticRefIds.Batch(SqliteDatabase.MaxBatchSize))
        {
            var placeholderIds = SqliteDatabase.MakeInPlaceholderParamIds(batch.Count);
            var sql = $@"
SELECT semref_id, range_json
FROM SemanticRefs WHERE semref_id IN ({SqliteDatabase.MakeInStatement(placeholderIds)})";
            _db.GetKeyValues(
                sql,
                cmd => cmd.AddPlaceholderParameters(placeholderIds, batch),
                (reader) => new KeyValuePair<int, TextRange>(
                    reader.GetInt32(0),
                    StorageSerializer.FromJson<TextRange>(reader.GetString(1))
                    ),
                rangeRows
               );
        }

        return ValueTask.FromResult((IList<TextRange>)rangeRows.GetValues(semanticRefIds));
    }

    public KnowledgeType GetKnowledgeType(int semanticRefId)
    {
        ArgumentVerify.ThrowIfLessThan(semanticRefId, 0, nameof(semanticRefId));

        return _db.Get(@"
SELECT knowledge_type
FROM SemanticRefs WHERE semref_id = @semref_id",
            (cmd) => cmd.AddParameter("semref_id", semanticRefId),
            (reader) => reader.GetString(0)
        );
    }

    public ValueTask<KnowledgeType> GetKnowledgeTypeAsync(int semanticRefId, CancellationToken cancellationToken = default)
    {
        return ValueTask.FromResult(GetKnowledgeType(semanticRefId));
    }

    public ValueTask<IList<KnowledgeType>> GetKnowledgeTypeAsync(IList<int> semanticRefIds, CancellationToken cancellationToken = default)
    {
        ArgumentVerify.ThrowIfNullOrEmpty(semanticRefIds, nameof(semanticRefIds));

        List<KnowledgeType> ranges = new(semanticRefIds.Count);
        foreach (var batch in semanticRefIds.Batch(SqliteDatabase.MaxBatchSize))
        {
            cancellationToken.ThrowIfCancellationRequested();

            var placeholderIds = SqliteDatabase.MakeInPlaceholderParamIds(batch.Count);
            var sql = $@"
SELECT knowledge_type
FROM SemanticRefs WHERE semref_id IN ({SqliteDatabase.MakeInStatement(placeholderIds)})
ORDER BY semref_id";
            var rows = _db.Enumerate(
                sql,
                cmd => cmd.AddPlaceholderParameters(placeholderIds, batch),
                (reader) => (KnowledgeType)reader.GetString(0)
               );
            ranges.AddRange(rows);
        }
        return ValueTask.FromResult((IList<KnowledgeType>)ranges);
    }


    public async IAsyncEnumerator<SemanticRef> GetAsyncEnumerator(CancellationToken cancellationToken = default)
    {
        using var cmd = _db.CreateCommand(@"
SELECT semref_id, range_json, knowledge_type, knowledge_json
FROM SemanticRefs
ORDER BY semref_id
");
        using var reader = cmd.ExecuteReader();
        while (await reader.ReadAsync(cancellationToken).ConfigureAwait(false))
        {
            SemanticRef semanticRef = ReadSemanticRef(reader);
            yield return semanticRef;
        }

    }

    public IList<SemanticRef> GetSlice(int startOrdinal, int endOrdinal)
    {
        ArgumentVerify.ThrowIfGreaterThan(startOrdinal, endOrdinal, nameof(startOrdinal));

        using var cmd = _db.CreateCommand(@"
SELECT semref_id, range_json, knowledge_type, knowledge_json
FROM SemanticRefs WHERE semref_id >= @start_id AND semref_id < @end_id
ORDER BY semref_id");
        cmd.AddParameter("@start_id", startOrdinal);
        cmd.AddParameter("@end_id", endOrdinal);

        using var reader = cmd.ExecuteReader();
        IList<SemanticRef> semanticRefList = reader.GetList(ReadSemanticRef);
        return semanticRefList;
    }


    public ValueTask<IList<SemanticRef>> GetSliceAsync(int startOrdinal, int endOrdinal, CancellationToken cancellationToken = default)
    {
        return ValueTask.FromResult(GetSlice(startOrdinal, endOrdinal));
    }

    public IList<SemanticRef> GetAll(KnowledgeType? kType = null)
    {
        return kType is null
            ? _db.GetList(
                "SELECT * FROM SemanticRefs",
                ReadSemanticRef
            )
            : _db.GetList(
@"SELECT * from SemanticRefs
WHERE knowledge_type = @kType",
            (cmd) => cmd.AddParameter("@kType", kType.ToString()),
            ReadSemanticRef
        );
    }

    public ValueTask<IList<SemanticRef>> GetAllAsync(KnowledgeType? kType = null, CancellationToken cancellationToken = default)
    {
        return ValueTask.FromResult(GetAll());
    }

    public IList<ScoredSemanticRefOrdinal> GetAllOrdinals(KnowledgeType? kType = null)
    {
        return kType is null
            ? _db.GetList(
                "SELECT semref_id, 1.0 FROM SemanticRefs",
                (reader) =>
                {
                    return new ScoredSemanticRefOrdinal
                    {
                        SemanticRefOrdinal = reader.GetInt32(0),
                        Score = reader.GetInt32(1)
                    };
                }
            )
            : _db.GetList(
@"SELECT semref_id, 1.0 from SemanticRefs
WHERE knowledge_type = @kType",
            (cmd) => cmd.AddParameter("@kType", kType.ToString()),
                (reader) =>
                {
                    return new ScoredSemanticRefOrdinal
                    {
                        SemanticRefOrdinal = reader.GetInt32(0),
                        Score = reader.GetInt32(1)
                    };
                }
        );
    }

    public ValueTask<IList<ScoredSemanticRefOrdinal>> GetAllOrdinalsAsync(KnowledgeType? kType = null, CancellationToken cancellationToken = default)
    {
        return ValueTask.FromResult(GetAllOrdinals(kType));
    }

    public event Action<BatchProgress> OnKnowledgeExtracted;

    public void NotifyKnowledgeProgress(BatchProgress progress)
    {
        OnKnowledgeExtracted.SafeInvoke(progress);
    }

    int GetNextSemanicRefId()
    {
        return GetCount();
    }

    SemanticRef ReadSemanticRef(SqliteDataReader reader)
    {
        var row = ReadSemanticRefRow(reader);
        return FromSemanticRefRow(row);
    }

    SemanticRefRow ToSemanticRefRow(SemanticRef semanticRef)
    {
        SemanticRefRow row = new();
        row.SemanticRefId = (semanticRef.SemanticRefOrdinal < 0)
            ? GetNextSemanicRefId()
            : semanticRef.SemanticRefOrdinal;
        row.RangeJson = StorageSerializer.ToJson(semanticRef.Range);
        row.KnowledgeType = semanticRef.KnowledgeType;
        row.KnowledgeJson = StorageSerializer.ToJson(semanticRef.Knowledge);
        return row;
    }

    SemanticRef FromSemanticRefRow(SemanticRefRow semanticRefRow)
    {
        SemanticRef semanticRef = new()
        {
            SemanticRefOrdinal = semanticRefRow.SemanticRefId,
            Range = StorageSerializer.FromJson<TextRange>(semanticRefRow.RangeJson),
            KnowledgeType = semanticRefRow.KnowledgeType,
            Knowledge = SemanticRef.Deserialize(semanticRefRow.KnowledgeJson, semanticRefRow.KnowledgeType)
        };

        return semanticRef;
    }

    SemanticRefRow ReadSemanticRefRow(SqliteDataReader reader)
    {
        SemanticRefRow row = new SemanticRefRow();

        int iCol = 0;
        row.SemanticRefId = reader.GetInt32(iCol++);
        row.RangeJson = reader.GetStringOrNull(iCol++);
        row.KnowledgeType = reader.GetStringOrNull(iCol++);
        row.KnowledgeJson = reader.GetStringOrNull(iCol++);

        return row;
    }

    void WriteSemanticRefRow(SqliteCommand cmd, SemanticRefRow semanticRefRow)
    {
        cmd.AddParameter("@semref_id", semanticRefRow.SemanticRefId);
        cmd.AddParameter("@range_json", semanticRefRow.RangeJson);
        cmd.AddParameter("@knowledge_type", semanticRefRow.KnowledgeType);
        cmd.AddParameter("@knowledge_json", semanticRefRow.KnowledgeJson);
    }
}

internal class SemanticRefRow
{
    public int SemanticRefId { get; set; }
    public string? RangeJson { get; set; }
    public string? KnowledgeType { get; set; }
    public string? KnowledgeJson { get; set; }
}
