using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using System.Threading;
using IntervalTree;
using Microsoft.Extensions.Logging;
using OmniSharp.Extensions.LanguageServer.Protocol;
using OmniSharp.Extensions.LanguageServer.Protocol.Models;

namespace Microsoft.Dafny.LanguageServer.Workspace;

record Reference(IHasNavigationToken Referer, IHasNavigationToken Referred);

public class SymbolTable {

  public static SymbolTable CreateFrom(ILogger logger, Program program, CancellationToken cancellationToken) {
    var visited = program.Visit(a => true, b => { });

    var references = visited.OfType<IHasReferences>().Where(v => !AutoGeneratedToken.Is(v.NavigationToken))
      .SelectMany(r => {
        var usages = r.GetReferences();
        if (usages == null) {
          logger.LogError($"Value of type {r.GetType()} returned a null for ${nameof(IHasReferences.GetReferences)}");
          usages = Array.Empty<IHasNavigationToken>();
        }
        return usages.Where(d => d != null).Select(declaration =>
          new Reference(r, declaration));
      }).ToList();

    var definitions = visited.OfType<IHasNavigationToken>();

    return new SymbolTable(references, definitions.ToList());
  }

  public static SymbolTable Empty() {
    return new SymbolTable();
  }

  private SymbolTable() {
    NodeToReferences = ImmutableDictionary<IHasNavigationToken, ISet<IHasNavigationToken>>.Empty;
    ReferenceToNode = ImmutableDictionary<IHasNavigationToken, IHasNavigationToken>.Empty;
    Symbols = ImmutableList<ISymbol>.Empty;
  }

  private SymbolTable(IReadOnlyList<Reference> references, IReadOnlyList<IHasNavigationToken> hasNavigationTokens) {
    var safeReferences = references.Where(k => k.Referer.NavigationToken.Uri != null).ToImmutableList();

    var referencesAsKeyValuePairs = safeReferences.Select(k => KeyValuePair.Create(k.Referer, k.Referred));
    ReferenceToNode = referencesAsKeyValuePairs.DistinctBy(pair => pair.Key).ToImmutableDictionary();

    NodeToReferences = safeReferences.GroupBy(u => u.Referred).ToImmutableDictionary(
      g => g.Key,
      g => (ISet<IHasNavigationToken>)g.Select(k => k.Referer).ToHashSet());

    var relevantDafnySymbolKinds = new HashSet<SymbolKind> {
      SymbolKind.Function,
      SymbolKind.Class,
      SymbolKind.Enum,
      SymbolKind.Method,
      SymbolKind.EnumMember,
      SymbolKind.Struct,
      SymbolKind.Interface,
      SymbolKind.Namespace,
    };
    // Since these definitions are checked for whether they
    // contain substrings when answering workspace/resolve queries,
    // performance can be improved by storing their names in a
    // data structure that makes this operation cheaper, such as
    // a suffix tree.
    Symbols = hasNavigationTokens.OfType<ISymbol>().
      Where(symbol => symbol.Kind.HasValue && relevantDafnySymbolKinds.Contains(symbol.Kind.Value)).
      ToImmutableList();

    var symbols = hasNavigationTokens
      .Concat(safeReferences.Select(r => r.Referer))
      .Concat(safeReferences.Select(r => r.Referred))
      .Where(s => !AutoGeneratedToken.Is(s.NavigationToken)).Distinct();
    var symbolsByFile = symbols.GroupBy(s => s.NavigationToken.Uri);
    foreach (var symbolsForFile in symbolsByFile) {
      if (symbolsForFile.Key == null) {
        continue;
      }
      var nodePositionsForFile = new IntervalTree<Position, IHasNavigationToken>();
      nodePositions.Add(symbolsForFile.Key, nodePositionsForFile);
      foreach (var symbolForFile in symbolsForFile) {
        var range = symbolForFile.NavigationToken.GetLspRange();
        nodePositionsForFile.Add(range.Start, range.End, symbolForFile);
      }
    }
  }

  private readonly Dictionary<Uri, IIntervalTree<Position, IHasNavigationToken>> nodePositions = new();

  /// <summary>
  /// Maps each symbol declaration to itself, and each symbol usage to the symbol's declaration.
  /// </summary>
  public ImmutableDictionary<IHasNavigationToken, IHasNavigationToken> ReferenceToNode { get; }

  /// <summary>
  /// Maps each symbol declaration to usages of the symbol, not including the declaration itself.
  /// </summary>
  public ImmutableDictionary<IHasNavigationToken, ISet<IHasNavigationToken>> NodeToReferences { get; }

  /// <summary>
  ///  A list of all definitions, such as methods, classes, functions, etc., used for workspace-wide symbol
  /// lookup.
  /// </summary>
  public ImmutableList<ISymbol> Symbols { get; }

  public IEnumerable<Location> GetReferences(Uri uri, Position position) {
    if (nodePositions.TryGetValue(uri, out var forFile)) {
      return forFile.Query(position).
        SelectMany(node => NodeToReferences.GetOrDefault(node, () => (ISet<IHasNavigationToken>)new HashSet<IHasNavigationToken>())).
        Select(u => new Location { Uri = u.NavigationToken.Filepath, Range = u.NavigationToken.GetLspRange() }).Distinct();
    }
    return Enumerable.Empty<Location>();
  }

  public Location? GetDeclaration(Uri uri, Position position) {
    var node = GetDeclarationNode(uri, position);
    return node == null ? null : NodeToLocation(node);
  }

  internal static Location? NodeToLocation(IHasNavigationToken node) {
    if (node.NavigationToken.Uri == null) {
      return null;
    }
    return new Location {
      Uri = DocumentUri.From(node.NavigationToken.Uri),
      Range = node.NavigationToken.GetLspRange()
    };
  }

  public IHasNavigationToken? GetDeclarationNode(Uri uri, Position position) {
    if (!nodePositions.TryGetValue(uri, out var forFile)) {
      return null;
    }
    return forFile.Query(position)
      .Select(node => node is ISymbol ? node : ReferenceToNode.GetValueOrDefault(node))
      .FirstOrDefault(x => x != null);
  }
}
