module Graph.DOT where

-- import Numeric.Probability
import Control.Applicative (Applicative (..))
import Control.Monad (Monad (..), foldM_, forM_, mapM_, when, zipWithM_)
import Control.Monad.Trans.Class qualified as MT
import Control.Monad.Trans.Reader qualified as MT
import Control.Monad.Trans.Writer.CPS qualified as MT
import Data.Bool (otherwise, (&&))
import Data.ByteString.Builder qualified as BS
import Data.ByteString.Short qualified as BSh
import Data.Eq (Eq (..))
import Data.Foldable (foldMap', toList)
import Data.Function (id, on, ($), (&), (.))
import Data.Functor ((<&>))
import Data.Int (Int)
import Data.List qualified as List
import Data.Map.Strict qualified as Map
import Data.Maybe (Maybe (..), fromMaybe)
import Data.Monoid (Monoid (..))
import Data.Ord (Ord (..))
import Data.Semigroup (Min (..), Semigroup (..))
import Data.Set (Set)
import Data.Set qualified as Set
import Data.String (IsString (..), String)
import Data.Text.Lazy qualified as LazyText
import Data.Text.Lazy.Encoding qualified as LazyText
import Data.Text.Short (ShortText)
import Data.Text.Short qualified as ShortText
import Data.Tree qualified as Tree
import Debug.Pretty.Simple (pTraceShow)
import GHC.Real (floor)
import GHC.Stack (HasCallStack)
import Numeric (showFFloat)
import Text.Pretty.Simple (pShow)
import Text.Printf (printf)
import Text.Show (Show (..))
import Prelude

type DOT m = MT.ReaderT BSh.ShortByteString (MT.WriterT BS.Builder m) ()

indexFrom1M_ :: Applicative m => [a] -> (a -> Int -> m b) -> m ()
indexFrom1M_ xs f = zipWithM_ f xs [1 :: Int ..]

runDOT :: Monad m => DOT m -> m BS.Builder
runDOT = MT.execWriterT . (`MT.runReaderT` {-indent-} "")

dotBlock :: Monad m => DOT m -> DOT m
dotBlock s = do
  dotLine "{"
  () <- MT.withReaderT ("  " <>) s
  dotLine "}"

dotLine :: Monad m => BS.Builder -> DOT m
dotLine s = do
  indent <- MT.ask
  MT.lift $ MT.tell $ BS.shortByteString indent <> s <> "\n"

dotComments :: Monad m => [BS.Builder] -> DOT m
dotComments cs = do
  dotLine "/*"
  forM_ cs dotLine
  dotLine "*/"

dotEdges :: Monad m => [BS.Builder] -> [(BS.Builder, Value)] -> DOT m
dotEdges names as = dotLine $ mconcat (List.intersperse " -> " names) <> builderAttrs as

data Value
  = ValueString String
  | ValueHTML BS.Builder
instance IsString Value where
  fromString = ValueString

dotNode :: Monad m => BS.Builder -> [(BS.Builder, Value)] -> DOT m
dotNode name as = dotLine $ name <> builderAttrs as

builderAttrs :: [(BS.Builder, Value)] -> BS.Builder
builderAttrs as
  | List.null as = ""
  | otherwise = "[" <> mconcat (List.intersperse "," attrs) <> "]"
  where
    attrs =
      [ k <> "=" <> vBS
      | (k, v) <- as
      , BS.toLazyByteString k /= ""
      , let vBS = case v of
              ValueString s -> s & builderQuotedString
              ValueHTML s -> BS.charUtf8 '<' <> s <> BS.charUtf8 '>'
      ]

builderQuotedShortText :: ShortText -> BS.Builder
builderQuotedShortText = builderQuotedString . ShortText.unpack

builderQuotedString :: String -> BS.Builder
builderQuotedString cs = BS.charUtf8 '"' <> foldMap' escape cs <> BS.charUtf8 '"'
  where
    escape '\\' = BS.charUtf8 '\\' <> BS.charUtf8 '\\'
    escape '\"' = BS.charUtf8 '\\' <> BS.charUtf8 '\"'
    escape c = BS.charUtf8 c