using System;
using System.Collections.Generic;
using System.Diagnostics.Contracts;
using System.Linq;

namespace Microsoft.Dafny;

public class InferDecreasesClause {
  private readonly ModuleResolver resolver;

  public InferDecreasesClause(ModuleResolver resolver) {
    this.resolver = resolver;
  }

  public void FillInDefaultDecreasesClauses(ModuleDefinition module) {

    Contract.Assert(Type.GetScope() != null);
    foreach (var clbl in ModuleDefinition.AllCallables(module.TopLevelDecls)) {
      ICallable m;
      string s;
      if (clbl is ExtremeLemma) {
        var prefixLemma = ((ExtremeLemma)clbl).PrefixLemma;
        m = prefixLemma;
        s = prefixLemma.Name + " ";
      } else {
        m = clbl;
        s = "";
      }

      var anyChangeToDecreases = FillInDefaultDecreases(m, true);

      if (anyChangeToDecreases || m.InferredDecreases || m is PrefixLemma) {
        bool showIt = false;
        if (m is Function) {
          // show the inferred decreases clause only if it will ever matter, i.e., if the function is recursive
          showIt = ((Function)m).IsRecursive;
        } else if (m is PrefixLemma) {
          // always show the decrease clause, since at the very least it will start with "_k", which the programmer did not write explicitly
          showIt = true;
        } else {
          showIt = ((Method)m).IsRecursive;
        }

        if (showIt) {
          s += "decreases " + Util.Comma(m.Decreases.Expressions, expr => Printer.ExprToString(resolver.Options, expr));
          // Note, in the following line, we use the location information for "clbl", not "m".  These
          // are the same, except in the case where "clbl" is a GreatestLemma and "m" is a prefix lemma.
          resolver.reporter.Info(MessageSource.Resolver, clbl.Origin, s);
        }
      }
    }
  }

  /// <summary>
  /// Return "true" if this routine makes any change to the decreases clause.  If the decreases clause
  /// starts off essentially empty and a default is provided, then clbl.InferredDecreases is also set
  /// to true.
  /// </summary>
  public bool FillInDefaultDecreases(ICallable clbl, bool addPrefixInCoClusters) {
    Contract.Requires(clbl != null);

    if (clbl is ExtremePredicate) {
      // extreme predicates don't have decreases clauses
      return false;
    }

    var anyChangeToDecreases = false;
    var decr = clbl.Decreases.Expressions;
    if (decr.Count == 0 || (clbl is PrefixLemma && decr.Count == 1)) {
      // The default for a function starts with the function's reads clause, if any
      if (clbl is Function) {
        var fn = (Function)clbl;
        if (fn.Reads.Expressions.Count != 0) {
          // start the default lexicographic tuple with the reads clause
          var r = FrameToObjectSet(fn.Reads.Expressions);
          decr.Add(AutoGeneratedExpression.Create(r));
          anyChangeToDecreases = true;
        }
      }

      if (clbl is Function || clbl is Method) {
        TopLevelDeclWithMembers enclosingType;
        MemberDecl originalMember;
        if (clbl is Function fc && !fc.IsStatic) {
          enclosingType = (TopLevelDeclWithMembers)fc.EnclosingClass;
          originalMember = fc.Original;
        } else if (clbl is Method mc && !mc.IsStatic) {
          enclosingType = (TopLevelDeclWithMembers)mc.EnclosingClass;
          originalMember = mc.Original;
        } else {
          enclosingType = null;
          originalMember = null;
        }

        if (enclosingType != null && originalMember == clbl) {
          var receiverType = ModuleResolver.GetThisType(clbl.Origin, enclosingType);
          if (receiverType.IsOrdered && !receiverType.IsRefType) {
            var th = new ThisExpr(new AutoGeneratedOrigin(clbl.Origin)) { Type = receiverType }; // resolve here
            decr.Add(AutoGeneratedExpression.Create(th));
            anyChangeToDecreases = true;
          }
        }
      }

      // Add one component for each parameter, unless the parameter's type is one that
      // doesn't appear useful to orderings.
      foreach (var p in clbl.Ins) {
        if (!(p is ImplicitFormal) && p.Type.IsOrdered) {
          var ie = new IdentifierExpr(new AutoGeneratedOrigin(p.Origin), p.Name);
          ie.Var = p;
          ie.Type = p.Type; // resolve it here
          decr.Add(AutoGeneratedExpression.Create(ie));
          anyChangeToDecreases = true;
        }
      }

      clbl.InferredDecreases = true; // this indicates that finding a default decreases clause was attempted
    }

    if (addPrefixInCoClusters && clbl is Function) {
      var fn = (Function)clbl;
      switch (fn.CoClusterTarget) {
        case Function.CoCallClusterInvolvement.None:
          break;
        case Function.CoCallClusterInvolvement.IsMutuallyRecursiveTarget:
          // prefix: decreases 0,
          clbl.Decreases.Expressions.Insert(0, Expression.CreateIntLiteral(fn.Origin, 0));
          anyChangeToDecreases = true;
          break;
        case Function.CoCallClusterInvolvement.CoRecursiveTargetAllTheWay:
          // prefix: decreases 1,
          clbl.Decreases.Expressions.Insert(0, Expression.CreateIntLiteral(fn.Origin, 1));
          anyChangeToDecreases = true;
          break;
        default:
          Contract.Assume(false); // unexpected case
          break;
      }
    }

    return anyChangeToDecreases;
  }

  public Expression FrameToObjectSet(List<FrameExpression> fexprs) {
    Contract.Requires(fexprs != null);
    Contract.Ensures(Contract.Result<Expression>() != null);

    List<Expression> sets = [];
    List<Expression> singletons = null;
    var idGen = new VerificationIdGenerator();
    // drop wildcards altogether in the following iterations
    foreach (FrameExpression fe in fexprs.Where(fe => fe.E is not WildcardExpr)) {
      Contract.Assert(fe != null);
      Expression e = new Cloner(false, true).CloneExpr(fe.E); // keep only fe.E, drop any fe.Field designation
      Contract.Assert(e.Type != null); // fe.E should have been resolved already, and the clone should still have that type
      var eType = e.Type.NormalizeToAncestorType();
      if (eType.IsRefType) {
        // e represents a singleton set
        if (singletons == null) {
          singletons = [];
        }
        singletons.Add(e);

      } else if (eType is SeqType or MultiSetType || LiteralExpr.IsEmptySet(e)) {
        // e represents a sequence or multiset
        var collectionType = (CollectionType)eType;
        var resolvedOpcode = collectionType.ResolvedOpcodeForIn;
        var boundedPool = collectionType.GetBoundedPool(e);

        // Add:  set x :: x in e
        var bvDecl = new BoundVar(e.Origin, idGen.FreshId("_s2s_"), collectionType.Arg.NormalizeExpand());
        var bv = new IdentifierExpr(e.Origin, bvDecl.Name) {
          Var = bvDecl,
          Type = bvDecl.Type
        };
        var bvInE = new BinaryExpr(e.Origin, BinaryExpr.Opcode.In, bv, e) {
          ResolvedOp = resolvedOpcode,
          Type = Type.Bool
        };
        var s = new SetComprehension(e.Origin, true, [bvDecl], bvInE, bv,
          new Attributes("trigger", [bvInE], null)) {
          Type = resolver.SystemModuleManager.ObjectSetType(),
          Bounds = [boundedPool]
        };
        sets.Add(s);

      } else {
        // e is already a set
        Contract.Assert(eType is SetType);
        sets.Add(e);
      }
    }

    if (singletons != null) {
      Expression display = new SetDisplayExpr(singletons[0].Origin, true, singletons) {
        Type = resolver.SystemModuleManager.ObjectSetType()
      };
      sets.Add(display);
    }

    if (sets.Count == 0) {
      var emptySet = new SetDisplayExpr(Token.NoToken, true, []) {
        Type = resolver.SystemModuleManager.ObjectSetType()
      };
      return emptySet;
    } else {
      Expression s = sets[0];
      for (var i = 1; i < sets.Count; i++) {
        var union = new BinaryExpr(s.Origin, BinaryExpr.Opcode.Add, s, sets[i]) {
          ResolvedOp = BinaryExpr.ResolvedOpcode.Union,
          Type = resolver.SystemModuleManager.ObjectSetType()
        };
        s = union;
      }

      return s;
    }
  }
}