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

namespace Microsoft.Dafny;

public class CalcStmt : Statement, ICloneable<CalcStmt>, ICanFormat {
  public abstract class CalcOp {
    /// <summary>
    /// Resulting operator "x op z" if "x this y" and "y other z".
    /// Returns null if this and other are incompatible.
    /// </summary>
    [System.Diagnostics.Contracts.Pure]
    public abstract CalcOp ResultOp(CalcOp other);

    /// <summary>
    /// Returns an expression "line0 this line1".
    /// </summary>
    [System.Diagnostics.Contracts.Pure]
    public abstract Expression StepExpr(Expression line0, Expression line1);
  }

  public class BinaryCalcOp : CalcOp {
    public readonly BinaryExpr.Opcode Op;

    [ContractInvariantMethod]
    void ObjectInvariant() {
      Contract.Invariant(ValidOp(Op));
    }

    /// <summary>
    /// Is op a valid calculation operator?
    /// </summary>
    [System.Diagnostics.Contracts.Pure]
    public static bool ValidOp(BinaryExpr.Opcode op) {
      return
        op == BinaryExpr.Opcode.Eq || op == BinaryExpr.Opcode.Neq
                                   || op == BinaryExpr.Opcode.Lt || op == BinaryExpr.Opcode.Le
                                   || op == BinaryExpr.Opcode.Gt || op == BinaryExpr.Opcode.Ge
                                   || LogicOp(op);
    }

    /// <summary>
    /// Is op a valid operator only for Boolean lines?
    /// </summary>
    [System.Diagnostics.Contracts.Pure]
    public static bool LogicOp(BinaryExpr.Opcode op) {
      return op == BinaryExpr.Opcode.Iff || op == BinaryExpr.Opcode.Imp || op == BinaryExpr.Opcode.Exp;
    }

    public BinaryCalcOp(BinaryExpr.Opcode op) {
      Contract.Requires(ValidOp(op));
      Op = op;
    }

    /// <summary>
    /// Does this subsume other (this . other == other . this == this)?
    /// </summary>
    private bool Subsumes(BinaryCalcOp other) {
      Contract.Requires(other != null);
      var op1 = Op;
      var op2 = other.Op;
      if (op1 == BinaryExpr.Opcode.Neq || op2 == BinaryExpr.Opcode.Neq) {
        return op2 == BinaryExpr.Opcode.Eq;
      }

      if (op1 == op2) {
        return true;
      }

      if (LogicOp(op1) || LogicOp(op2)) {
        return op2 == BinaryExpr.Opcode.Eq ||
               (op1 == BinaryExpr.Opcode.Imp && op2 == BinaryExpr.Opcode.Iff) ||
               (op1 == BinaryExpr.Opcode.Exp && op2 == BinaryExpr.Opcode.Iff) ||
               (op1 == BinaryExpr.Opcode.Eq && op2 == BinaryExpr.Opcode.Iff);
      }

      return op2 == BinaryExpr.Opcode.Eq ||
             (op1 == BinaryExpr.Opcode.Lt && op2 == BinaryExpr.Opcode.Le) ||
             (op1 == BinaryExpr.Opcode.Gt && op2 == BinaryExpr.Opcode.Ge);
    }

    public override CalcOp ResultOp(CalcOp other) {
      if (other is BinaryCalcOp) {
        var o = (BinaryCalcOp)other;
        if (Subsumes(o)) {
          return this;
        } else if (o.Subsumes(this)) {
          return other;
        }
        return null;
      } else if (other is TernaryCalcOp) {
        return other.ResultOp(this);
      } else {
        Contract.Assert(false);
        throw new cce.UnreachableException();
      }
    }

    public override Expression StepExpr(Expression line0, Expression line1) {
      if (Op == BinaryExpr.Opcode.Exp) {
        // The order of operands is reversed so that it can be turned into implication during resolution
        return new BinaryExpr(new AutoGeneratedOrigin(line0.Origin), Op, line1, line0);
      } else {
        return new BinaryExpr(new AutoGeneratedOrigin(line0.Origin), Op, line0, line1);
      }
    }

    public override string ToString() {
      return BinaryExpr.OpcodeString(Op);
    }

  }

  public class TernaryCalcOp : CalcOp {
    public readonly Expression Index; // the only allowed ternary operator is ==#, so we only store the index

    [ContractInvariantMethod]
    void ObjectInvariant() {
      Contract.Invariant(Index != null);
    }

    public TernaryCalcOp(Expression idx) {
      Contract.Requires(idx != null);
      Index = idx;
    }

    public override CalcOp ResultOp(CalcOp other) {
      if (other is BinaryCalcOp) {
        if (((BinaryCalcOp)other).Op == BinaryExpr.Opcode.Eq) {
          return this;
        }
        return null;
      } else if (other is TernaryCalcOp) {
        var a = Index;
        var b = ((TernaryCalcOp)other).Index;
        var minIndex = new ITEExpr(a.Origin, false, new BinaryExpr(a.Origin, BinaryExpr.Opcode.Le, a, b), a, b);
        return new TernaryCalcOp(minIndex); // ToDo: if we could compare expressions for syntactic equality, we could use this here to optimize
      } else {
        Contract.Assert(false);
        throw new cce.UnreachableException();
      }
    }

    public override Expression StepExpr(Expression line0, Expression line1) {
      return new TernaryExpr(new AutoGeneratedOrigin(line0.Origin), TernaryExpr.Opcode.PrefixEqOp, Index, line0, line1);
    }

    public override string ToString() {
      return "==#";
    }

  }

  /// <summary>
  /// This method infers a default operator to be used between the steps.
  /// Usually, we'd use == as the default operator.  However, if the calculation
  /// begins or ends with a boolean literal, then we can do better by selecting ==>
  /// or <==.  Also, if the calculation begins or ends with an empty set, then we can
  /// do better by selecting <= or >=.
  /// Note, these alternative operators are chosen only if they don't clash with something
  /// supplied by the user.
  /// If the rules come up with a good inferred default operator, then that default operator
  /// is returned; otherwise, null is returned.
  /// </summary>
  [CanBeNull]
  public CalcOp GetInferredDefaultOp() {
    CalcOp alternativeOp = null;
    if (Lines.Count == 0) {
      return null;
    }

    if (Expression.IsBoolLiteral(Lines.First(), out var firstOperatorIsBoolLiteral)) {
      alternativeOp = new BinaryCalcOp(firstOperatorIsBoolLiteral ? BinaryExpr.Opcode.Imp : BinaryExpr.Opcode.Exp);
    } else if (Expression.IsBoolLiteral(Lines.Last(), out var lastOperatorIsBoolLiteral)) {
      alternativeOp = new BinaryCalcOp(lastOperatorIsBoolLiteral ? BinaryExpr.Opcode.Exp : BinaryExpr.Opcode.Imp);
    } else if (Expression.IsEmptySetOrMultiset(Lines.First())) {
      alternativeOp = new BinaryCalcOp(BinaryExpr.Opcode.Ge);
    } else if (Expression.IsEmptySetOrMultiset(Lines.Last())) {
      alternativeOp = new BinaryCalcOp(BinaryExpr.Opcode.Le);
    } else {
      return null;
    }

    // Check that the alternative operator is compatible with anything supplied by the user.
    var resultOp = alternativeOp;
    foreach (var stepOp in StepOps.Where(stepOp => stepOp != null)) {
      resultOp = resultOp.ResultOp(stepOp);
      if (resultOp == null) {
        // no go
        return null;
      }
    }
    return alternativeOp;
  }

  public readonly List<Expression> Lines;    // Last line is dummy, in order to form a proper step with the dangling hint
  public readonly List<BlockStmt> Hints;     // Hints[i] comes after line i; block statement is used as a container for multiple sub-hints
  public readonly CalcOp UserSuppliedOp;     // may be null, if omitted by the user
  public readonly List<CalcOp/*?*/> StepOps; // StepOps[i] comes after line i
  [FilledInDuringResolution]
  public CalcOp Op;                          // main operator of the calculation (either UserSuppliedOp or (after resolution) an inferred CalcOp)
  [FilledInDuringResolution] public readonly List<Expression> Steps;    // expressions li op l<i + 1> (last step is dummy)
  [FilledInDuringResolution] public Expression Result;                  // expression l0 ResultOp ln

  public static readonly CalcOp DefaultOp = new BinaryCalcOp(BinaryExpr.Opcode.Eq);

  public override IEnumerable<INode> Children => Steps.Concat(Result != null ? [Result] : new Node[] { }).Concat(Hints);
  public override IEnumerable<INode> PreResolveChildren => Lines.Take(Lines.Count > 0 ? Lines.Count - 1 : 0).Concat<Node>(Hints.Where(hintBatch => hintBatch.Body.Count() != 0));

  [ContractInvariantMethod]
  void ObjectInvariant() {
    Contract.Invariant(Lines != null);
    Contract.Invariant(cce.NonNullElements(Lines));
    Contract.Invariant(Hints != null);
    Contract.Invariant(cce.NonNullElements(Hints));
    Contract.Invariant(StepOps != null);
    Contract.Invariant(Steps != null);
    Contract.Invariant(cce.NonNullElements(Steps));
    Contract.Invariant(Hints.Count == Math.Max(Lines.Count - 1, 0));
    Contract.Invariant(StepOps.Count == Hints.Count);
  }

  public CalcStmt(IOrigin origin, CalcOp userSuppliedOp, List<Expression> lines, List<BlockStmt> hints, List<CalcOp/*?*/> stepOps, Attributes attrs)
    : base(origin) {
    Contract.Requires(origin != null);
    Contract.Requires(lines != null);
    Contract.Requires(hints != null);
    Contract.Requires(stepOps != null);
    Contract.Requires(cce.NonNullElements(lines));
    Contract.Requires(cce.NonNullElements(hints));
    Contract.Requires(hints.Count == Math.Max(lines.Count - 1, 0));
    Contract.Requires(stepOps.Count == hints.Count);
    UserSuppliedOp = userSuppliedOp;
    Lines = lines;
    Hints = hints;
    Steps = [];
    StepOps = stepOps;
    Result = null;
    Attributes = attrs;
  }

  public CalcStmt Clone(Cloner cloner) {
    return new CalcStmt(cloner, this);
  }

  public CalcStmt(Cloner cloner, CalcStmt original) : base(cloner, original) {
    // calc statements have the unusual property that the last line is duplicated.  If that is the case (which
    // we expect it to be here), we share the clone of that line as well.
    var lineCount = original.Lines.Count;
    var lines = new List<Expression>(lineCount);
    for (int i = 0; i < lineCount; i++) {
      lines.Add(i == lineCount - 1 && 2 <= lineCount && original.Lines[i] == original.Lines[i - 1] ? lines[i - 1] : cloner.CloneExpr(original.Lines[i]));
    }
    UserSuppliedOp = cloner.CloneCalcOp(original.UserSuppliedOp);
    Lines = lines;
    StepOps = original.StepOps.ConvertAll(cloner.CloneCalcOp);
    Hints = original.Hints.ConvertAll(cloner.CloneBlockStmt);

    if (cloner.CloneResolvedFields) {
      Steps = original.Steps.Select(cloner.CloneExpr).ToList();
      Result = cloner.CloneExpr(original.Result);
      Op = original.Op;
    } else {
      Steps = [];
    }
  }

  public override IEnumerable<Statement> SubStatements {
    get {
      foreach (var h in Hints) {
        yield return h;
      }
    }
  }
  public override IEnumerable<Expression> SpecificationSubExpressions {
    get {
      foreach (var e in base.SpecificationSubExpressions) { yield return e; }
      foreach (var e in Attributes.SubExpressions(Attributes)) { yield return e; }

      for (int i = 0; i < Lines.Count - 1; i++) {  // note, we skip the duplicated line at the end
        yield return Lines[i];
      }
      foreach (var calcop in AllCalcOps) {
        if (calcop is TernaryCalcOp o3) {
          yield return o3.Index;
        }
      }

      if (Result != null) {
        yield return Result;
      }
    }
  }

  IEnumerable<CalcOp> AllCalcOps {
    get {
      if (UserSuppliedOp != null) {
        yield return UserSuppliedOp;
      }
      foreach (var stepop in StepOps) {
        if (stepop != null) {
          yield return stepop;
        }
      }
    }
  }

  /// <summary>
  /// Left-hand side of a step expression.
  /// Note that Lhs(op.StepExpr(line0, line1)) != line0 when op is <==.
  /// </summary>
  public static Expression Lhs(Expression step) {
    Contract.Requires(step is BinaryExpr || step is TernaryExpr);
    if (step is BinaryExpr) {
      return ((BinaryExpr)step).E0;
    } else {
      return ((TernaryExpr)step).E1;
    }
  }

  /// <summary>
  /// Right-hand side of a step expression.
  /// Note that Rhs(op.StepExpr(line0, line1)) != line1 when op is REVERSE-IMPLICATION.
  /// </summary>
  public static Expression Rhs(Expression step) {
    Contract.Requires(step is BinaryExpr || step is TernaryExpr);
    if (step is BinaryExpr) {
      return ((BinaryExpr)step).E1;
    } else {
      return ((TernaryExpr)step).E2;
    }
  }

  public bool SetIndent(int indentBefore, TokenNewIndentCollector formatter) {
    var inCalc = false;
    var inOrdinal = false;
    var innerCalcIndent = indentBefore + formatter.SpaceTab;
    var extraHintIndent = 0;
    var ownedTokens = OwnedTokens;
    // First phase: We get the alignment
    foreach (var token in ownedTokens) {
      if (formatter.SetIndentLabelTokens(token, indentBefore)) {
        continue;
      }
      switch (token.val) {
        case "calc":
        case ";":
        case "}": {
            break;
          }
        case "{": {
            inCalc = true;
            break;
          }
        default: {
            if (inCalc) {
              if (token.val == "[") {
                inOrdinal = true;
              }
              if (token.val == "]") {
                inOrdinal = false;
              }
              if (!TokenNewIndentCollector.IsFollowedByNewline(token) &&
                  (token.val != "==" || token.Next.val != "#") &&
                  token.val != "#" &&
                  !inOrdinal) {
                if (token.Next.val != "{") {
                  formatter.SetIndentations(token, inline: indentBefore);
                  innerCalcIndent = Math.Max(innerCalcIndent, formatter.GetRightAlignIndentAfter(token, indentBefore));
                } else {// It's an hint! If there is no comment and no newline between them, we align the hints as well.
                  if ((token.TrailingTrivia + token.Next.LeadingTrivia).Trim() == "" &&
                      token.line == token.Next.line) {
                    extraHintIndent = Math.Max(extraHintIndent, formatter.GetRightAlignIndentAfter(token, indentBefore) - (indentBefore + formatter.SpaceTab));
                  }
                }
              }
            }

            break;
          }
      }
    }

    inCalc = false;
    foreach (var token in OwnedTokens) {
      switch (token.val) {
        case "calc": {
            break;
          }
        case "{": {
            formatter.SetIndentations(token, indentBefore, indentBefore, innerCalcIndent);
            inCalc = true;
            break;
          }
        case "}": {
            formatter.SetIndentations(token, innerCalcIndent, indentBefore, indentBefore);
            break;
          }
        case ";": {
            formatter.SetDelimiterInsideIndentedRegions(token, indentBefore);
            break;
          }
        default: {
            // It has to be an operator
            if (inCalc) {
              formatter.SetIndentations(token, innerCalcIndent, indentBefore, innerCalcIndent);
            }

            break;
          }
      }
    }

    foreach (var hint in Hints) {
      // This block
      if (hint.Origin.pos != hint.EndToken.pos) {
        foreach (var hintStep in hint.Body) {
          formatter.SetOpeningIndentedRegion(hintStep.StartToken, indentBefore + formatter.SpaceTab + extraHintIndent);
        }
      }
    }

    foreach (var expression in Lines) {
      formatter.SetIndentations(expression.StartToken, innerCalcIndent, innerCalcIndent, innerCalcIndent);
    }

    return true;
  }

  public override void ResolveGhostness(ModuleResolver resolver, ErrorReporter reporter, bool mustBeErasable,
    ICodeContext codeContext, string proofContext,
    bool allowAssumptionVariables, bool inConstructorInitializationPhase) {
    IsGhost = true;
    foreach (var hint in Hints) {
      hint.ResolveGhostness(resolver, reporter, true, codeContext, "a hint", allowAssumptionVariables,
        inConstructorInitializationPhase);
    }
  }
}