// Copyright (c) 2023, the Dart project authors.  Please see the AUTHORS file
// for details. All rights reserved. Use of this source code is governed by a
// BSD-style license that can be found in the LICENSE file.
import '../piece/piece.dart';
import '../profile.dart';
import 'code.dart';
import 'code_writer.dart';
import 'solution_cache.dart';

/// A single possible set of formatting choices.
///
/// Each solution binds some number of [Piece]s in the piece tree to [State]s.
/// (Any pieces whose states are not bound are treated as having a default
/// unsplit state.)
///
/// Given that set of states, we can create a [CodeWriter] and give that to all
/// of the pieces in the tree so they can format themselves. That in turn
/// yields a total number of overflow characters, cost, and formatted output,
/// which are all stored here.
final class Solution implements Comparable<Solution> {
  /// The states that pieces have been bound to.
  ///
  /// Note that order that keys are inserted into this map is significant. When
  /// ordering solutions, we use the order that pieces are bound in here to
  /// break ties between solutions that otherwise have the same cost and
  /// overflow.
  final Map<Piece, State> _pieceStates;

  /// The set of states that pieces are allowed to be in without violating
  /// constraints of already bound pieces.
  ///
  /// Each key is a constrained piece and the values are the remaining states
  /// that the piece may take which aren't known to violate existing
  /// constraints. If a piece is not in this map, then there are no constraints
  /// on it.
  final Map<Piece, List<State>> _allowedStates;

  /// The amount of penalties applied based on the chosen line splits.
  int get cost => _cost + _subtreeCost;

  /// The cost of this solution based on pieces it has bound to states itself,
  /// excluding pieces from separately formatted subtrees.
  int _cost;

  /// The cost of this solution from branches of the piece tree that were
  /// separately formatted and merged in using [mergeSubtree()].
  ///
  /// We track this separately so that when expanding a solution, we don't
  /// double count the cost of separately formatted branches.
  int _subtreeCost = 0;

  /// The formatted code.
  GroupCode get code => _code;
  late final GroupCode _code;

  /// False if this Solution contains a newline where one is prohibited.
  ///
  /// An invalid solution may have no overflow characters and the lowest score,
  /// but is still considered worse than any other valid solution.
  bool get isValid => _isValid;
  bool _isValid = true;

  /// Whether the solution contains an invalid newline and the piece that
  /// prohibits the newline is bound in this solution.
  ///
  /// When this is `true`, it means this solution and every solution that could
  /// be derived from it is invalid so the whole solution tree can be discarded.
  bool _isDeadEnd = false;

  /// The total number of characters that do not fit inside the page width.
  int get overflow => _overflow;
  int _overflow = 0;

  /// The unsolved piece in this solution that should be expanded next to
  /// produce new more refined solutions, if there is one.
  ///
  /// The tree of possible solutions is combinatorial in the number of pieces
  /// and exponential in the number of states those pieces can take. We can't
  /// afford to brute force explore the whole tree, even with the optimization
  /// that we stop as soon as we find a solution with no overflow.
  ///
  /// Most possible solutions add unnecessary splits in regions of the code
  /// that already fit within the page width. Exploring those is wasted time.
  /// To avoid that, we rely on a couple of insights:
  ///
  /// First, the solver treats any piece with an unselected state as being
  /// unsplit. This means that refining a solution always takes a piece that is
  /// unsplit and makes it split more. That monotonically increases the cost,
  /// but may help fit the solution inside the page.
  ///
  /// Therefore, we don't want to select states for most pieces. Only pieces
  /// that need to split in order to find a solution that fits in the page
  /// width or that are necessary because the unsplit state is invalid. (The
  /// latter usually means a line comment or statement occurs inside the piece.)
  ///
  /// So we skip past any pieces that aren't on overflowing lines or on lines
  /// whose newline led to an invalid solution. Further, it's also the case
  /// that splitting earlier pieces will often reshuffle the formatting of much
  /// of the code following it.
  ///
  /// Thus we only worry about unsolved pieces on the *first* problematic line
  /// when expanding. If selecting states for those pieces still doesn't help,
  /// the solver will work its way through later pieces from those subsequent
  /// partial solutions.
  ///
  /// This lets us efficiently skip through almost all of the pieces that don't
  /// need to be touched in order to find a valid solution.
  ///
  /// If this is empty, then there are no further solutions to generate from
  /// this one. It's either a dead end or a winner.
  late final List<Piece> _expandPieces;

  /// Creates a new [Solution] with no pieces set to any state (which
  /// implicitly means they have state [State.unsplit] unless they're pinned to
  /// another state).
  factory Solution(
    SolutionCache cache,
    Piece root, {
    required int pageWidth,
    required int leadingIndent,
    required int subsequentIndent,
    State? rootState,
  }) {
    var solution = Solution._(root, 0, {}, {}, rootState);
    solution._format(cache, root, pageWidth, leadingIndent, subsequentIndent);
    return solution;
  }

  Solution._(
    Piece root,
    this._cost,
    this._pieceStates,
    this._allowedStates, [
    State? rootState,
  ]) {
    Profile.count('create Solution');

    // If we're formatting a subtree of a larger Piece tree that binds [root]
    // to [rootState], then bind it in this solution too.
    if (rootState != null) _bind(root, rootState);
  }

  /// Attempt to eagerly bind [piece] to a state given that it must fit within
  /// [pageWidth] (which is the overall page width minus any leading indentation
  /// in the solution where this is called).
  ///
  /// If it can, binds the piece to that state in this solution and returns
  /// `true`. Otherwise returns `false`.
  bool tryBindByPageWidth(Piece piece, int pageWidth) {
    if (piece.fixedStateForPageWidth(pageWidth) case var state?) {
      _bind(piece, state);
      return true;
    }

    return false;
  }

  /// The state that [piece] is pinned to or that this solution selects.
  ///
  /// If no state has been selected, defaults to the first state.
  State pieceState(Piece piece) => pieceStateIfBound(piece) ?? State.unsplit;

  /// The state that [piece] is pinned to or that this solution selects.
  ///
  /// If no state has been selected, returns `null`.
  State? pieceStateIfBound(Piece piece) =>
      piece.pinnedState ?? _pieceStates[piece];

  /// Whether [piece] has been bound to a state in this set (or is pinned).
  bool isBound(Piece piece) =>
      piece.pinnedState != null || _pieceStates.containsKey(piece);

  /// Increases the total overflow for this solution by [overflow].
  ///
  /// This should only be called by [CodeWriter].
  void addOverflow(int overflow) {
    _overflow += overflow;
  }

  /// Apply the overflow, cost, and bound states from [subtreeSolution] to this
  /// solution.
  ///
  /// This is called when a subtree of a Piece tree is solved separately and
  /// the resulting solution is being merged with this one.
  void mergeSubtree(Solution subtreeSolution) {
    _overflow += subtreeSolution._overflow;
    _subtreeCost += subtreeSolution.cost;
  }

  /// Mark this solution as having a newline where none is permitted by [piece]
  /// and is thus not a valid solution.
  ///
  /// This should only be called by [CodeWriter].
  void invalidate(Piece piece) {
    // Don't invalidate if the piece is pinned and can't do anything. This only
    // comes into play when a mandatory newline inside a string interpolation
    // causes some surrounding interpolation expression to try to split but the
    // expression is aready pinned to be unsplit to prevent newlines inside
    // interpolation.
    // TODO(rnystrom): If we come up with a cleaner way to model how newlines
    // inside interpolation are handled, we can get rid of this. See the TODO
    // comment in visitInterpolationString().
    if (piece.pinnedState != null) return;

    _isValid = false;

    // TODO(rnystrom): It would be good to mark a solution as a dead end here
    // if possible. It used to be safe to do that whenever [piece] was bound
    // because any newline at all from any child would reliably always
    // invalidate that piece, so if bound, it must be a dead end.
    //
    // Now that we invalidate based on allowed shapes, and the shape of a piece
    // can depend on its children (transitively), it's not clear when we can
    // safely consider a solution a dead end when gets invalidated.
  }

  /// Derives new potential solutions from this one by binding [_expandPieces]
  /// to all of their possible states.
  ///
  /// If there is no potential piece to expand, or all attempts to expand it
  /// fail, returns an empty list.
  List<Solution> expand(
    SolutionCache cache,
    Piece root, {
    required int pageWidth,
    required int leadingIndent,
    required int subsequentIndent,
  }) {
    // If there is no piece that we can expand on this solution, it's a dead
    // end (or a winner).
    if (_expandPieces.isEmpty) return const [];

    var solutions = <Solution>[];
    for (var i = 0; i < _expandPieces.length; i++) {
      // For each non-default state that the expanding piece can be in, create
      // a new solution that inherits all of the bindings of this one, and binds
      // the expanding piece to that state (along with any further pieces
      // constrained by that one).
      var expandPiece = _expandPieces[i];
      for (var state
          in _allowedStates[expandPiece] ?? expandPiece.additionalStates) {
        var expanded = Solution._(
          root,
          _cost,
          {..._pieceStates},
          {..._allowedStates},
        );

        // Bind all preceding expand pieces to their unsplit state. Their
        // other states have already been expanded by earlier iterations of
        // the outer for loop.
        var valid = true;
        for (var j = 0; j < i; j++) {
          expanded._bind(_expandPieces[j], State.unsplit);

          if (expanded._isDeadEnd) {
            valid = false;
            break;
          }
        }

        // Discard the solution if we hit a constraint violation.
        if (!valid) continue;

        expanded._bind(expandPiece, state);

        // Discard the solution if we hit a constraint violation.
        if (!expanded._isDeadEnd) {
          expanded._format(
            cache,
            root,
            pageWidth,
            leadingIndent,
            subsequentIndent,
          );

          // TODO(rnystrom): These come mostly (entirely?) from hard newlines
          // in sequences, comments, and multiline strings. It should be
          // possible to handle those during piece construction too. If we do,
          // remove this check.
          // We may not detect some newline violations until formatting.
          if (!expanded._isDeadEnd) solutions.add(expanded);
        }
      }
    }

    return solutions;
  }

  /// Compares two solutions where a more desirable solution comes first.
  ///
  /// For performance, we want to stop checking solutions as soon as we find
  /// the best one. Best means the fewest overflow characters and the lowest
  /// code.
  @override
  int compareTo(Solution other) {
    // Even though overflow is "worse" than cost, we order in terms of cost
    // because a solution with overflow may lead to a low-cost solution without
    // overflow, so we want to explore in cost order.
    if (cost != other.cost) return cost.compareTo(other.cost);

    if (overflow != other.overflow) return overflow.compareTo(other.overflow);

    // If all else is equal, prefer lower states in earlier bound pieces.
    for (var piece in _pieceStates.keys) {
      var thisState = pieceState(piece);
      var otherState = other.pieceState(piece);
      if (thisState != otherState) return thisState.compareTo(otherState);
    }

    return 0;
  }

  @override
  String toString() {
    var states = [
      for (var MapEntry(key: piece, value: state) in _pieceStates.entries)
        if (piece.additionalStates.isNotEmpty && piece.pinnedState == null)
          '$piece$state',
    ];

    return [
      '\$$cost ($_cost + $_subtreeCost)',
      if (overflow > 0) '($overflow over)',
      if (!isValid) '(invalid)',
      states.join(' '),
    ].join(' ').trim();
  }

  /// Run a [CodeWriter] on this solution to produce the final formatted output
  /// and calculate the overflow and expand pieces.
  void _format(
    SolutionCache cache,
    Piece root,
    int pageWidth,
    int leadingIndent,
    int subsequentIndent,
  ) {
    var writer = CodeWriter(
      pageWidth,
      leadingIndent,
      subsequentIndent,
      cache,
      this,
    );

    writer.format(root);

    var (code, expandPieces) = writer.finish();
    _code = code;
    _expandPieces = expandPieces;
  }

  /// Attempts to add a binding from [piece] to [state] to the solution, and
  /// then adds any further bindings from constraints that [piece] applies to
  /// its children, recursively.
  ///
  /// This may invalidate the solution if [piece] is already bound to a
  /// different [state], or if any constrained pieces are bound to different
  /// states.
  ///
  /// If successful, adds the cost required to bind [piece] to [state] (along
  /// with any other applied constrained pieces). Otherwise, marks the solution
  /// as a dead end.
  void _bind(Piece piece, State state) {
    // If we've already failed from a previous violation, early out.
    if (_isDeadEnd) return;

    // Apply the new binding if it doesn't conflict with an existing one.
    switch (pieceStateIfBound(piece)) {
      case null:
        // Binding a unbound piece to a state.
        _cost += piece.stateCost(state);
        _pieceStates[piece] = state;

        // This piece may in turn place further constraints on others.
        piece.applyConstraints(state, _bind);

        // If this piece's state prevents some of its children from having
        // newlines, then further constrain those children.
        if (!_isDeadEnd) {
          piece.forEachChild((child) {
            // Stop as soon as we fail.
            if (_isDeadEnd) return;

            // If the child can't have newlines in any shape, then constrain it.
            // TODO(rnystrom): We can probably constrain this more specifically
            // by asking what shapes the child might have.
            var allowedShapes = piece.allowedChildShapes(state, child);
            if (allowedShapes.length == 1 &&
                allowedShapes.contains(Shape.inline)) {
              _constrainOffspring(child);
            }
          });
        }

      case var alreadyBound when alreadyBound != state:
        // Already bound to a different state, so there's a conflict.
        _isDeadEnd = true;
        _isValid = false;

      default:
        break; // Already bound to the same state, so nothing to do.
    }
  }

  /// For [piece] and its transitive offspring subtree, eliminate any state that
  /// will always produce a newline since that state is not permitted because
  /// the parent of [piece] doesn't allow [piece] to have any newlines.
  void _constrainOffspring(Piece piece) {
    for (var offspring in piece.statefulOffspring) {
      if (_isDeadEnd) break;

      if (pieceStateIfBound(offspring) case var boundState?) {
        // This offspring is already pinned or bound to a state. If that
        // state will emit newlines, then this solution is invalid.
        if (offspring.containsNewline(boundState)) {
          _isDeadEnd = true;
          _isValid = false;
        }
      } else if (!_allowedStates.containsKey(offspring)) {
        // If we get here, the offspring isn't bound to a state and we haven't
        // already constrained it. Eliminate any of its states that will emit
        // newlines.
        var allowedUnsplit = !offspring.containsNewline(State.unsplit);

        var states = offspring.additionalStates;
        var remainingStates = <State>[];
        for (var state in states) {
          if (!offspring.containsNewline(state)) {
            remainingStates.add(state);
          }
        }

        if (!allowedUnsplit && remainingStates.isEmpty) {
          // There is no state this child can take that won't emit newlines,
          // and it's not allowed to, so this solution is bad.
          _isDeadEnd = true;
          _isValid = false;
        } else if (remainingStates.isEmpty) {
          // The only valid state is unsplit so bind it to that.
          _bind(offspring, State.unsplit);
        } else if (!allowedUnsplit && remainingStates.length == 1) {
          // There's only one valid state, so bind it to that.
          _bind(offspring, remainingStates.first);
        } else if (remainingStates.length < states.length) {
          // There are some constrained states, so keep the remaining ones.
          _allowedStates[offspring] = remainingStates;
        }
      }
    }
  }
}
