{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Language.Haskell.Brittany.Internal.Utils where

import qualified Data.ByteString as B
import qualified Data.Coerce
import Data.Data
import Data.Generics.Aliases
import qualified Data.Generics.Uniplate.Direct as Uniplate
import DataTreePrint
import qualified GHC.Data.FastString as GHC
import qualified GHC.OldList as List
import           GHC.Types.Name.Occurrence            (OccName, occNameString)
import qualified GHC.Types.Name.Reader as RdrName (rdrNameOcc)
import qualified GHC.Types.SrcLoc as GHC
import qualified GHC.Utils.Outputable as GHC
import qualified GHC.Parser.Annotation as GHC
-- import qualified Language.Haskell.GHC.ExactPrint.Utils as ExactPrint.Utils
import qualified Text.PrettyPrint as PP
import qualified Data.Semigroup as Semigroup
import qualified System.IO.Unsafe as Unsafe

import Language.Haskell.Brittany.Internal.Prelude
import Language.Haskell.Brittany.Internal.Components.BriDoc
import Language.Haskell.Brittany.Internal.Types
import Language.Haskell.Brittany.Internal.Config.Types



parDoc :: String -> PP.Doc
parDoc = PP.fsep . fmap PP.text . List.words

parDocW :: [String] -> PP.Doc
parDocW = PP.fsep . fmap PP.text . List.words . List.unwords


showSDoc_ :: GHC.SDoc -> String
showSDoc_ = GHC.showSDocUnsafe

showOutputable :: (GHC.Outputable a) => a -> String
showOutputable = GHC.showPprUnsafe

fromMaybeIdentity :: Identity a -> Maybe a -> Identity a
fromMaybeIdentity x y = Data.Coerce.coerce $ fromMaybe (Data.Coerce.coerce x) y

fromOptionIdentity :: Identity a -> Maybe a -> Identity a
fromOptionIdentity x y =
  Data.Coerce.coerce $ fromMaybe (Data.Coerce.coerce x) y

-- maximum monoid over N+0
-- or more than N, because Num is allowed.
newtype Max a = Max { getMax :: a }
  deriving (Eq, Ord, Show, Num)

instance (Num a, Ord a) => Semigroup (Max a) where
  (<>) = Data.Coerce.coerce (max :: a -> a -> a)

instance (Num a, Ord a) => Monoid (Max a) where
  mempty = Max 0
  mappend = (<>)

newtype ShowIsId = ShowIsId String deriving Data

instance Show ShowIsId where
  show (ShowIsId x) = x

data A x = A ShowIsId x
  deriving Data

data DeltaComment = DeltaComment GHC.DeltaPos GHC.EpaCommentTok
  deriving Data

customLayouterF :: LayouterF
customLayouterF layoutF =
  DataToLayouter
    $ f
    `extQ` internalLayouterShowIsId
    `extQ` internalLayouterFastString
    `extQ` internalLayouterBytestring
    `extQ` internalLayouterOccName
    `extQ` internalLayouterSrcSpan
    `extQ` internalLayouterRdrName
    `extQ` realSrcSpan
    -- `extQ` deltaComment
    -- `extQ` anchored
    -- `ext1Q` srcSpanAnn
    -- `ext2Q` located
 where
  DataToLayouter f = defaultLayouterF layoutF
  realSrcSpan :: GHC.RealSrcSpan -> NodeLayouter
  realSrcSpan span = internalLayouterSimple (show span)
  -- anchored :: (Data b) => GHC.GenLocated GHC.Anchor b -> NodeLayouter
  -- anchored (GHC.L (GHC.Anchor _ op) a) = f $ GHC.L op a

customLayouterNoSrcSpansF :: LayouterF
customLayouterNoSrcSpansF layoutF =
  DataToLayouter
    $ f
    `extQ` internalLayouterShowIsId
    `extQ` internalLayouterFastString
    `extQ` internalLayouterBytestring
    `extQ` internalLayouterOccName
    `extQ` internalLayouterSrcSpan
    `extQ` internalLayouterRdrName
    `extQ` realSrcSpan
    `extQ` anchored
    `ext1Q` srcSpanAnn
    -- `ext2Q` located
 where
  DataToLayouter f = defaultLayouterF layoutF
  realSrcSpan :: GHC.RealSrcSpan -> NodeLayouter
  realSrcSpan span = internalLayouterSimple (show span)
  -- anchored :: (Data b) => GHC.GenLocated GHC.Anchor b -> NodeLayouter
  -- anchored (GHC.L (GHC.Anchor _ op) a) = f $ GHC.L op a
  anchored :: GHC.Anchor -> NodeLayouter
  anchored (GHC.Anchor _ op) = f op
  srcSpanAnn :: forall a . Data a => GHC.SrcSpanAnn' a -> NodeLayouter
  srcSpanAnn (GHC.SrcSpanAnn ann _loc) = f ann
  -- deltaComment :: GHC.LEpaComment -> NodeLayouter
  -- deltaComment (GHC.L anchor (GHC.EpaComment token prior)) =
  --   f (GHC.L (ExactPrint.Utils.ss2deltaEnd prior (GHC.anchor anchor)) token)
  -- located :: (Data b, Data loc) => GHC.GenLocated loc b -> NodeLayouter
  -- located (GHC.L ss a) = runDataToLayouter layoutF $ A annStr a
  --  where
  --   annStr = case cast ss of
  --     Just (s :: GHC.SrcSpan) ->
  --       ShowIsId $ show (ExactPrint.Utils.getAnnotationEP (GHC.L s a) anns)
  --     Nothing -> ShowIsId "nnnnnnnn"

customLayouterNoAnnsF :: LayouterF
customLayouterNoAnnsF layoutF =
  DataToLayouter
    $ f
    `extQ` internalLayouterShowIsId
    `extQ` internalLayouterFastString
    `extQ` internalLayouterBytestring
    `extQ` internalLayouterOccName
    `extQ` internalLayouterSrcSpan
    `extQ` internalLayouterRdrName
    `extQ` realSrcSpan
    `extQ` realSrcLoc
    `ext2Q` located
    `extQ` lepaComment
 where
  DataToLayouter f = defaultLayouterF layoutF
  located :: (Data b) => GHC.GenLocated loc b -> NodeLayouter
  located (GHC.L _ss a) = runDataToLayouter layoutF a
  lepaComment :: GHC.LEpaComment -> NodeLayouter
  lepaComment (GHC.L anchor (GHC.EpaComment token _)) = f
    ( token
    , GHC.srcSpanStartLine (GHC.anchor anchor)
    , GHC.srcSpanStartCol (GHC.anchor anchor)
    , GHC.srcSpanEndLine (GHC.anchor anchor)
    , GHC.srcSpanEndCol (GHC.anchor anchor)
    )
  realSrcSpan :: GHC.RealSrcSpan -> NodeLayouter
  realSrcSpan span = internalLayouterSimple
    (show
      ( GHC.srcSpanStartLine span
      , GHC.srcSpanStartCol span
      , GHC.srcSpanEndLine span
      , GHC.srcSpanEndCol span
      )
    )
  realSrcLoc :: GHC.RealSrcLoc -> NodeLayouter
  realSrcLoc loc =
    internalLayouterSimple (show (GHC.srcLocLine loc, GHC.srcLocCol loc))

internalLayouterSimple :: String -> NodeLayouter
internalLayouterSimple s = NodeLayouter (length s) False (const $ PP.text s)
internalLayouterShowIsId :: ShowIsId -> NodeLayouter
internalLayouterShowIsId (ShowIsId s) =
  NodeLayouter (length s + 2) True $ \case
    Left True -> PP.parens $ PP.text s
    Left False -> PP.text s
    Right _ -> PP.text s
internalLayouterFastString :: GHC.FastString -> NodeLayouter
internalLayouterFastString =
  internalLayouterSimple . ("{FastString: " ++) . (++ "}") . show :: GHC.FastString
    -> NodeLayouter
internalLayouterBytestring :: B.ByteString -> NodeLayouter
internalLayouterBytestring =
  internalLayouterSimple . show :: B.ByteString -> NodeLayouter
internalLayouterSrcSpan :: GHC.SrcSpan -> NodeLayouter
internalLayouterSrcSpan ss =
    internalLayouterSimple
             -- - $ "{"++ showSDoc_ (GHC.ppr ss)++"}"
      $ "{"
      ++ showOutputable ss
      ++ "}"
internalLayouterOccName :: OccName -> NodeLayouter
internalLayouterOccName =
  internalLayouterSimple . ("{OccName: " ++) . (++ "}") . occNameString
internalLayouterRdrName :: RdrName -> NodeLayouter
internalLayouterRdrName =
  internalLayouterSimple . ("{RdrName: " ++) . (++ "}") . occNameString . RdrName.rdrNameOcc


-- displayBriDocTree :: BriDoc -> PP.Doc
-- displayBriDocTree = \case
--   BDWrapAnnKey annKey doc -> def "BDWrapAnnKey"
--                            $ PP.text (show annKey)
--                          $+$ displayBriDocTree doc
--   BDEmpty         -> PP.text "BDEmpty"
--   BDLit t         -> def "BDLit" $ PP.text (show t)
--   BDSeq list      -> def "BDSeq" $ displayList list
--   BDCols sig list -> def "BDCols" $ PP.text (show sig)
--                                 $+$ displayList list
--   BDSeparator     -> PP.text "BDSeparator"
--   BDPar rol indent lines -> def "BDPar" $ displayBriDocTree rol
--                                       $+$ PP.text (show indent)
--                                       $+$ displayList lines
--   BDAlt alts      -> def "BDAlt" $ displayList alts
--   BDExternal ast _t -> def "BDExternal" (astToDoc ast)
--   BDSpecialPostCommentLoc _ -> PP.text "BDSpecialPostCommentLoc"
--  where
--   def x r = PP.text x $+$ PP.nest 2 r
--   displayList :: [BriDoc] -> PP.Doc
--   displayList [] = PP.text "[]"
--   displayList (x:xr) = PP.cat $ PP.text "[" <+> displayBriDocTree x
--                               : [PP.text "," <+> displayBriDocTree t | t<-xr]
--                              ++ [PP.text "]"]

-- displayBriDocSimpleTree :: BriDocSimple -> PP.Doc
-- displayBriDocSimpleTree = \case
--   BDSWrapAnnKey annKey doc -> def "BDSWrapAnnKey"
--                            $ PP.text (show annKey)
--                          $+$ displayBriDocSimpleTree doc
--   BDSLit t         -> def "BDSLit" $ PP.text (show t)
--   BDSSeq list      -> def "BDSSeq" $ displayList list
--   BDSCols sig list -> def "BDSCols" $ PP.text (show sig)
--                                 $+$ displayList list
--   BDSSeparator     -> PP.text "BDSSeparator"
--   BDSPar rol indent lines -> def "BDSPar" $ displayBriDocSimpleTree rol
--                                       $+$ PP.text (show indent)
--                                       $+$ displayList lines
--   BDSExternal annKey _subKeys _t -> def "BDSExternal" (PP.text $ show annKey)
--   BDSSpecialPostCommentLoc _ -> PP.text "BDSSpecialPostCommentLoc"
--  where
--   def x r = PP.text x $+$ PP.nest 2 r
--   displayList :: [BriDocSimple] -> PP.Doc
--   displayList [] = PP.text "[]"
--   displayList (x:xr) = PP.cat $ PP.text "[" <+> displayBriDocSimpleTree x
--                               : [PP.text "," <+> displayBriDocSimpleTree t | t<-xr]
--                              ++ [PP.text "]"]

-- traceWhen
--   :: (Show a)
--   => String
--   -> Bool
--   -> a
--   -> IO ()
-- traceWhen s accessor val = do
--   TraceFunc f <- mAsk
--   whenM (mAsk <&> _conf_debug .> accessor .> confUnpack) $ do
--     Unsafe.unsafePerformIO $ do
--       f ("---- " ++ s ++ " ----\n" ++ show val)
--       pure $ pure ()


astToDoc :: Data ast => ast -> PP.Doc
astToDoc ast = printTreeWithCustom 160 customLayouterNoAnnsF ast

briDocToDoc :: BriDoc -> PP.Doc
briDocToDoc = astToDoc . removeAnnotations
 where
  removeAnnotations = Uniplate.transform $ \case
    -- BDFlushCommentsPrior _ x -> x
    -- BDFlushCommentsPost _ _ x -> x
    -- BDQueueComments _ x -> x
    x -> x

briDocToDocWithAnns :: BriDoc -> PP.Doc
briDocToDocWithAnns = astToDoc

-- annsDoc :: ExactPrint.Types.Anns -> PP.Doc
-- annsDoc =
--   printTreeWithCustom 100 customLayouterNoAnnsF . fmap (ShowIsId . show)

breakEither :: (a -> Either b c) -> [a] -> ([b], [c])
breakEither _ [] = ([], [])
breakEither fn (a1 : aR) = case fn a1 of
  Left b -> (b : bs, cs)
  Right c -> (bs, c : cs)
  where (bs, cs) = breakEither fn aR

spanMaybe :: (a -> Maybe b) -> [a] -> ([b], [a])
spanMaybe f (x1 : xR) | Just y <- f x1 = (y : ys, xs)
  where (ys, xs) = spanMaybe f xR
spanMaybe _ xs = ([], xs)

data FirstLastView a
  = FirstLastEmpty
  | FirstLastSingleton a
  | FirstLast a [a] a

splitFirstLast :: [a] -> FirstLastView a
splitFirstLast [] = FirstLastEmpty
splitFirstLast [x] = FirstLastSingleton x
splitFirstLast (x1 : xr) = FirstLast x1 (List.init xr) (List.last xr)

-- | similar to List.lines, but treating the case of final newline character
-- in such a manner that this function is the inverse of @intercalate "\n"@.
lines' :: String -> [String]
lines' s = case break (== '\n') s of
  (s1, []) -> [s1]
  (s1, [_]) -> [s1, ""]
  (s1, (_ : r)) -> s1 : lines' r

-- absurdExt :: HsExtension.NoExtCon -> a
-- absurdExt = HsExtension.noExtCon

traceIfDumpConf
  :: (MonadMultiReader TraceFunc m, MonadMultiReader Config m, Show a)
  => String
  -> (DebugConfig -> Identity (Semigroup.Last Bool))
  -> a
  -> m ()
traceIfDumpConf s accessor val = do
  TraceFunc f <- mAsk
  whenM (mAsk <&> _conf_debug .> accessor .> confUnpack) $ do
    Unsafe.unsafePerformIO $ do
      f ("---- " ++ s ++ " ----\n" ++ show val)
      pure $ pure ()

mergeOn :: Ord b => (a -> b) -> [a] -> [a] -> [a]
mergeOn _f xs [] = xs
mergeOn _f [] ys = ys
mergeOn f xs@(x:xr) ys@(y:yr)
    | f x <= f y = x : mergeOn f xr ys
    | otherwise = y : mergeOn f xs yr