module Language.PureScript.Sugar.Operators.Common where

import Prelude

import Control.Monad.State (guard, join)
import Control.Monad.Except (MonadError(..))

import Data.Either (rights)
import Data.Functor.Identity (Identity)
import Data.List (sortOn)
import Data.Maybe (mapMaybe, fromJust)
import Data.List.NonEmpty qualified as NEL
import Data.Map qualified as M

import Text.Parsec qualified as P
import Text.Parsec.Pos qualified as P
import Text.Parsec.Expr qualified as P

import Language.PureScript.AST (Associativity(..), ErrorMessageHint(..), SourceSpan)
import Language.PureScript.Crash (internalError)
import Language.PureScript.Errors (ErrorMessage(..), MultipleErrors(..), SimpleErrorMessage(..))
import Language.PureScript.Names (OpName, Qualified, eraseOpName)

type Chain a = [Either a a]

type FromOp nameType a = a -> Maybe (SourceSpan, Qualified (OpName nameType))
type Reapply nameType a = SourceSpan -> Qualified (OpName nameType) -> a -> a -> a

toAssoc :: Associativity -> P.Assoc
toAssoc Infixl = P.AssocLeft
toAssoc Infixr = P.AssocRight
toAssoc Infix  = P.AssocNone

token :: (P.Stream s Identity t) => (t -> Maybe a) -> P.Parsec s u a
token = P.token (const "") (const (P.initialPos ""))

parseValue :: P.Parsec (Chain a) () a
parseValue = token (either Just (const Nothing)) P.<?> "expression"

parseOp
  :: FromOp nameType a
  -> P.Parsec (Chain a) () (SourceSpan, Qualified (OpName nameType))
parseOp fromOp = token (either (const Nothing) fromOp) P.<?> "operator"

matchOp
  :: FromOp nameType a
  -> Qualified (OpName nameType)
  -> P.Parsec (Chain a) () SourceSpan
matchOp fromOp op = do
  (ss, ident) <- parseOp fromOp
  guard $ ident == op
  pure ss

opTable
  :: [[(Qualified (OpName nameType), Associativity)]]
  -> FromOp nameType a
  -> Reapply nameType a
  -> [[P.Operator (Chain a) () Identity a]]
opTable ops fromOp reapply =
  map (map (\(name, a) -> P.Infix (P.try (matchOp fromOp name) >>= \ss -> return (reapply ss name)) (toAssoc a))) ops

matchOperators
  :: forall m a nameType
   . Show a
  => MonadError MultipleErrors m
  => (a -> Bool)
  -> (a -> Maybe (a, a, a))
  -> FromOp nameType a
  -> Reapply nameType a
  -> ([[P.Operator (Chain a) () Identity a]] -> P.OperatorTable (Chain a) () Identity a)
  -> [[(Qualified (OpName nameType), Associativity)]]
  -> a
  -> m a
matchOperators isBinOp extractOp fromOp reapply modOpTable ops = parseChains
  where
  parseChains :: a -> m a
  parseChains ty
    | True <- isBinOp ty = bracketChain (extendChain ty)
    | otherwise = pure ty
  extendChain :: a -> Chain a
  extendChain ty
    | Just (op, l, r) <- extractOp ty = Left l : Right op : extendChain r
    | otherwise = [Left ty]
  bracketChain :: Chain a -> m a
  bracketChain chain =
    case P.parse opParser "operator expression" chain of
      Right a -> pure a
      Left _ -> throwError . MultipleErrors $ mkErrors chain
  opParser :: P.Parsec (Chain a) () a
  opParser = P.buildExpressionParser (modOpTable (opTable ops fromOp reapply)) parseValue <* P.eof

  -- Generating a good error message involves a bit of work here, as the parser
  -- can't provide one for us.
  --
  -- We examine the expression chain, plucking out the operators and then
  -- grouping them by shared precedence, then if any of the following conditions
  -- are met, we have something to report:
  --   1. any of the groups have mixed associativity
  --   2. there is more than one occurrence of a non-associative operator in a
  --      precedence group
  mkErrors :: Chain a -> [ErrorMessage]
  mkErrors chain =
    let
      opInfo :: M.Map (Qualified (OpName nameType)) (Integer, Associativity)
      opInfo = M.fromList $ concatMap (\(n, o) -> map (\(name, assoc) -> (name, (n, assoc))) o) (zip [0..] ops)
      opPrec :: Qualified (OpName nameType) -> Integer
      opPrec = fst . fromJust . flip M.lookup opInfo
      opAssoc :: Qualified (OpName nameType) -> Associativity
      opAssoc = snd . fromJust . flip M.lookup opInfo
      chainOpSpans :: M.Map (Qualified (OpName nameType)) (NEL.NonEmpty SourceSpan)
      chainOpSpans = foldr (\(ss, name) -> M.alter (Just . maybe (pure ss) (NEL.cons ss)) name) M.empty . mapMaybe fromOp $ rights chain
      opUsages :: Qualified (OpName nameType) -> Int
      opUsages = maybe 0 NEL.length . flip M.lookup chainOpSpans
      precGrouped :: [NEL.NonEmpty (Qualified (OpName nameType))]
      precGrouped = NEL.groupWith opPrec . sortOn opPrec $ M.keys chainOpSpans
      assocGrouped :: [NEL.NonEmpty (NEL.NonEmpty (Qualified (OpName nameType)))]
      assocGrouped = fmap (NEL.groupWith1 opAssoc . NEL.sortWith opAssoc) precGrouped
      mixedAssoc :: [NEL.NonEmpty (Qualified (OpName nameType))]
      mixedAssoc = fmap join . filter (\precGroup -> NEL.length precGroup > 1) $ assocGrouped
      nonAssoc :: [NEL.NonEmpty (Qualified (OpName nameType))]
      nonAssoc = NEL.filter (\assocGroup -> opAssoc (NEL.head assocGroup) == Infix && sum (fmap opUsages assocGroup) > 1) =<< assocGrouped
    in
      if null (nonAssoc ++ mixedAssoc)
        then internalError "matchOperators: cannot reorder operators"
        else
          map
            (\grp ->
              mkPositionedError chainOpSpans grp
                (MixedAssociativityError (fmap (\name -> (eraseOpName <$> name, opAssoc name)) grp)))
            mixedAssoc
          ++ map
            (\grp ->
              mkPositionedError chainOpSpans grp
                (NonAssociativeError (fmap (fmap eraseOpName) grp)))
            nonAssoc

  mkPositionedError
    :: M.Map (Qualified (OpName nameType)) (NEL.NonEmpty SourceSpan)
    -> NEL.NonEmpty (Qualified (OpName nameType))
    -> SimpleErrorMessage
    -> ErrorMessage
  mkPositionedError chainOpSpans grp =
    ErrorMessage
      [PositionedError (fromJust . flip M.lookup chainOpSpans =<< grp)]