// Copyright (C) 2015-2026 The Neo Project.
//
// StatePlugin.cs file belongs to the neo project and is free
// software distributed under the MIT software license, see the
// accompanying file LICENSE in the main directory of the
// repository or http://www.opensource.org/licenses/mit-license.php
// for more details.
//
// Redistribution and use in source and binary forms with or without
// modifications are permitted.

using Akka.Actor;
using Neo.ConsoleService;
using Neo.Cryptography.MPTTrie;
using Neo.Extensions;
using Neo.Json;
using Neo.Network.P2P.Payloads;
using Neo.Persistence;
using Neo.Plugins.RpcServer;
using Neo.Plugins.StateService.Storage;
using Neo.Plugins.StateService.Verification;
using Neo.SmartContract;
using Neo.SmartContract.Native;
using Neo.Wallets;
using System.Buffers.Binary;
using static Neo.Ledger.Blockchain;

namespace Neo.Plugins.StateService;

public class StatePlugin : Plugin
{
    public const string StatePayloadCategory = "StateService";
    public override string Name => "StateService";
    public override string Description => "Enables MPT for the node";
    public override string ConfigFile => System.IO.Path.Combine(RootPath, "StateService.json");

    protected override UnhandledExceptionPolicy ExceptionPolicy => StateServiceSettings.Default.ExceptionPolicy;

    internal IActorRef Store;
    internal IActorRef Verifier;

    private static NeoSystem _system;

    internal static NeoSystem NeoSystem => _system;

    private IWalletProvider walletProvider;

    public StatePlugin()
    {
        Committing += Blockchain_Committing_Handler;
        Committed += Blockchain_Committed_Handler;
    }

    protected override void Configure()
    {
        StateServiceSettings.Load(GetConfiguration());
    }

    protected override void OnSystemLoaded(NeoSystem system)
    {
        if (system.Settings.Network != StateServiceSettings.Default.Network) return;
        _system = system;
        Store = _system.ActorSystem.ActorOf(StateStore.Props(this, string.Format(StateServiceSettings.Default.Path, system.Settings.Network.ToString("X8"))));
        _system.ServiceAdded += NeoSystem_ServiceAdded_Handler;
        RpcServerPlugin.RegisterMethods(this, StateServiceSettings.Default.Network);
    }

    void NeoSystem_ServiceAdded_Handler(object sender, object service)
    {
        if (service is IWalletProvider)
        {
            walletProvider = service as IWalletProvider;
            _system.ServiceAdded -= NeoSystem_ServiceAdded_Handler;
            if (StateServiceSettings.Default.AutoVerify)
            {
                walletProvider.WalletChanged += IWalletProvider_WalletChanged_Handler;
            }
        }
    }

    void IWalletProvider_WalletChanged_Handler(object sender, Wallet wallet)
    {
        walletProvider.WalletChanged -= IWalletProvider_WalletChanged_Handler;
        Start(wallet);
    }

    public override void Dispose()
    {
        Committing -= Blockchain_Committing_Handler;
        Committed -= Blockchain_Committed_Handler;
        if (Store is not null) _system.EnsureStopped(Store);
        if (Verifier is not null) _system.EnsureStopped(Verifier);
        base.Dispose();
    }

    void Blockchain_Committing_Handler(NeoSystem system, Block block, DataCache snapshot,
        IReadOnlyList<ApplicationExecuted> applicationExecutedList)
    {
        if (system.Settings.Network != StateServiceSettings.Default.Network) return;
        StateStore.Singleton.UpdateLocalStateRootSnapshot(block.Index,
            snapshot.GetChangeSet()
                .Where(p => p.Value.State != TrackState.None && p.Key.Id != NativeContract.Ledger.Id)
                .ToList());
    }

    void Blockchain_Committed_Handler(NeoSystem system, Block block)
    {
        if (system.Settings.Network != StateServiceSettings.Default.Network) return;
        StateStore.Singleton.UpdateLocalStateRoot(block.Index);
    }

    private void CheckNetwork()
    {
        var network = StateServiceSettings.Default.Network;
        if (_system is null || _system.Settings.Network != network)
            throw new InvalidOperationException($"Network doesn't match: {_system?.Settings.Network} != {network}");
    }

    [ConsoleCommand("start states", Category = "StateService", Description = "Start as a state verifier if wallet is open")]
    private void OnStartVerifyingState()
    {
        CheckNetwork();
        Start(walletProvider.GetWallet());
    }

    public void Start(Wallet wallet)
    {
        if (Verifier is not null)
        {
            ConsoleHelper.Warning("Already started!");
            return;
        }
        if (wallet is null)
        {
            Verifier = null;
            ConsoleHelper.Warning("Please open wallet first!");
            return;
        }
        Verifier = _system.ActorSystem.ActorOf(VerificationService.Props(wallet));
    }

    [ConsoleCommand("state root", Category = "StateService", Description = "Get state root by index")]
    private void OnGetStateRoot(uint index)
    {
        CheckNetwork();

        using var snapshot = StateStore.Singleton.GetSnapshot();
        var stateRoot = snapshot.GetStateRoot(index);
        if (stateRoot is null)
            ConsoleHelper.Warning("Unknown state root");
        else
            ConsoleHelper.Info(stateRoot.ToJson().ToString());
    }

    [ConsoleCommand("state height", Category = "StateService", Description = "Get current state root index")]
    private void OnGetStateHeight()
    {
        CheckNetwork();

        ConsoleHelper.Info("LocalRootIndex: ",
            $"{StateStore.Singleton.LocalRootIndex}",
            " ValidatedRootIndex: ",
            $"{StateStore.Singleton.ValidatedRootIndex}");
    }

    [ConsoleCommand("get proof", Category = "StateService", Description = "Get proof of key and contract hash")]
    private void OnGetProof(UInt256 rootHash, UInt160 scriptHash, string key)
    {
        if (_system is null || _system.Settings.Network != StateServiceSettings.Default.Network)
            throw new InvalidOperationException("Network doesn't match");

        try
        {
            ConsoleHelper.Info("Proof: ", GetProof(rootHash, scriptHash, Convert.FromBase64String(key)));
        }
        catch (RpcException e)
        {
            ConsoleHelper.Error(e.Message);
        }
    }

    [ConsoleCommand("verify proof", Category = "StateService", Description = "Verify proof, return value if successed")]
    private void OnVerifyProof(UInt256 rootHash, string proof)
    {
        try
        {
            ConsoleHelper.Info("Verify Result: ", VerifyProof(rootHash, Convert.FromBase64String(proof)));
        }
        catch (RpcException e)
        {
            ConsoleHelper.Error(e.Message);
        }
    }

    [RpcMethod]
    public JToken GetStateRoot(uint index)
    {
        using var snapshot = StateStore.Singleton.GetSnapshot();
        var stateRoot = snapshot.GetStateRoot(index).NotNull_Or(RpcError.UnknownStateRoot);
        return stateRoot.ToJson();
    }

    private string GetProof(Trie trie, int contractId, byte[] key)
    {
        var skey = new StorageKey()
        {
            Id = contractId,
            Key = key,
        };
        return GetProof(trie, skey);
    }

    private string GetProof(Trie trie, StorageKey skey)
    {
        trie.TryGetProof(skey.ToArray(), out var proof).True_Or(RpcError.UnknownStorageItem);
        using var ms = new MemoryStream();
        using var writer = new BinaryWriter(ms, Utility.StrictUTF8);

        writer.WriteVarBytes(skey.ToArray());
        writer.WriteVarInt(proof.Count);
        foreach (var item in proof)
        {
            writer.WriteVarBytes(item);
        }
        writer.Flush();

        return Convert.ToBase64String(ms.ToArray());
    }

    private string GetProof(UInt256 rootHash, UInt160 scriptHash, byte[] key)
    {
        CheckRootHash(rootHash);

        using var store = StateStore.Singleton.GetStoreSnapshot();
        var trie = new Trie(store, rootHash);
        var contract = GetHistoricalContractState(trie, scriptHash).NotNull_Or(RpcError.UnknownContract);
        return GetProof(trie, contract.Id, key);
    }

    [RpcMethod]
    public JToken GetProof(UInt256 rootHash, UInt160 scriptHash, string key)
    {
        var keyBytes = Result.Ok_Or(() => Convert.FromBase64String(key), RpcError.InvalidParams.WithData($"Invalid key: {key}"));
        return GetProof(rootHash, scriptHash, keyBytes);
    }

    private string VerifyProof(UInt256 rootHash, byte[] proof)
    {
        var proofs = new HashSet<byte[]>();

        using var ms = new MemoryStream(proof, false);
        using var reader = new BinaryReader(ms, Utility.StrictUTF8);

        var key = reader.ReadVarBytes(Node.MaxKeyLength);
        var count = reader.ReadVarInt(byte.MaxValue);
        for (ulong i = 0; i < count; i++)
        {
            proofs.Add(reader.ReadVarBytes());
        }

        var value = Trie.VerifyProof(rootHash, key, proofs).NotNull_Or(RpcError.InvalidProof);
        return Convert.ToBase64String(value);
    }

    [RpcMethod]
    public JToken VerifyProof(UInt256 rootHash, string proof)
    {
        var proofBytes = Result.Ok_Or(
            () => Convert.FromBase64String(proof), RpcError.InvalidParams.WithData($"Invalid proof: {proof}"));
        return VerifyProof(rootHash, proofBytes);
    }

    [RpcMethod]
    public JToken GetStateHeight()
    {
        return new JObject()
        {
            ["localrootindex"] = StateStore.Singleton.LocalRootIndex,
            ["validatedrootindex"] = StateStore.Singleton.ValidatedRootIndex,
        };
    }

    private ContractState GetHistoricalContractState(Trie trie, UInt160 scriptHash)
    {
        const byte prefix = 8;
        var skey = new KeyBuilder(NativeContract.ContractManagement.Id, prefix).Add(scriptHash);
        return trie.TryGetValue(skey.ToArray(), out var value)
            ? value.AsSerializable<StorageItem>().GetInteroperable<ContractState>()
            : null;
    }

    private StorageKey ParseStorageKey(byte[] data)
    {
        return new() { Id = BinaryPrimitives.ReadInt32LittleEndian(data), Key = data.AsMemory(sizeof(int)) };
    }

    private void CheckRootHash(UInt256 rootHash)
    {
        var fullState = StateServiceSettings.Default.FullState;
        var current = StateStore.Singleton.CurrentLocalRootHash;
        (!fullState && current != rootHash)
            .False_Or(RpcError.UnsupportedState.WithData($"fullState:{fullState},current:{current},rootHash:{rootHash}"));
    }

    [RpcMethod]
    public JToken FindStates(UInt256 rootHash, UInt160 scriptHash, byte[] prefix, byte[] key = null, int count = 0)
    {
        CheckRootHash(rootHash);

        key ??= [];
        count = count <= 0 ? StateServiceSettings.Default.MaxFindResultItems : count;
        count = Math.Min(count, StateServiceSettings.Default.MaxFindResultItems);

        using var store = StateStore.Singleton.GetStoreSnapshot();
        var trie = new Trie(store, rootHash);
        var contract = GetHistoricalContractState(trie, scriptHash).NotNull_Or(RpcError.UnknownContract);
        var pkey = new StorageKey() { Id = contract.Id, Key = prefix };
        var fkey = new StorageKey() { Id = pkey.Id, Key = key };

        var json = new JObject();
        var jarr = new JArray();
        int i = 0;
        foreach (var (ikey, ivalue) in trie.Find(pkey.ToArray(), 0 < key.Length ? fkey.ToArray() : null))
        {
            if (count < i) break;
            if (i < count)
            {
                jarr.Add(new JObject()
                {
                    ["key"] = Convert.ToBase64String(ParseStorageKey(ikey.ToArray()).Key.Span),
                    ["value"] = Convert.ToBase64String(ivalue.Span),
                });
            }
            i++;
        }
        if (0 < jarr.Count)
        {
            json["firstProof"] = GetProof(trie, contract.Id, Convert.FromBase64String(jarr.First()["key"].AsString()));
        }
        if (1 < jarr.Count)
        {
            json["lastProof"] = GetProof(trie, contract.Id, Convert.FromBase64String(jarr.Last()["key"].AsString()));
        }
        json["truncated"] = count < i;
        json["results"] = jarr;
        return json;
    }

    [RpcMethod]
    public JToken GetState(UInt256 rootHash, UInt160 scriptHash, byte[] key)
    {
        CheckRootHash(rootHash);

        using var store = StateStore.Singleton.GetStoreSnapshot();
        var trie = new Trie(store, rootHash);
        var contract = GetHistoricalContractState(trie, scriptHash).NotNull_Or(RpcError.UnknownContract);
        var skey = new StorageKey()
        {
            Id = contract.Id,
            Key = key,
        };
        return Convert.ToBase64String(trie[skey.ToArray()]);
    }
}
