{-# options_ghc -Wno-orphans -Wno-inline-rule-shadowing #-}

module Common (
    module Common
  , module Control.Applicative
  , module Control.Monad
  , module Control.Monad.IO.Class
  , module Control.Monad.Reader.Class
  , module Control.Monad.State.Class
  , module Data.Bits
  , module Data.Flat
  , module Data.Foldable
  , module Data.IORef
  , module GHC.Exts
  , module GHC.Word
  , module Lens.Micro
  , module Lens.Micro.TH
  , module Utils.List
  , module Utils.Pair
  , IO(..)
  , FP.strToUtf8
  , FP.utf8ToStr
  , ppShow
  , pPrint
  , trace
  , traceM
  , traceShow
  , traceShowM
  ) where

import Control.Applicative hiding (Const(..))
import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.Reader.Class
import Control.Monad.State.Class
import Data.Bits
import Data.ByteString (ByteString)
import Data.ByteString qualified as B
import Data.ByteString.Internal qualified as B
import Data.Flat hiding (size)
import Data.Foldable
import Data.Hashable
import Data.IORef
import Data.List (intercalate)
import Data.Ref.F qualified as RF
import Data.Ref.L qualified as RL
import Debug.Trace
import FlatParse.Basic qualified as FP
import GHC.Exts hiding (lazy, toList, List, BCO, mkApUpd0#, newBCO#)
import GHC.ForeignPtr
import GHC.IO
import GHC.Word
import Lens.Micro
import Lens.Micro.TH
import Text.Show.Pretty (ppShow, pPrint)

import Utils.List
import Utils.Pair


-- debug printing, toggled by DEBUG cabal flag
----------------------------------------------------------------------------------------------------

-- TODO: more sophistication!
--   toggleable hierarchical debug printing, at least Agda-level!
--   - meta-annotations on expressions to toggle debug printing, benchmarking etc
--     expression-level commands that have side effect when elab gets to them
--     like DEBUG[foo,bar,baz]( expr )

#ifdef DEBUG
import GHC.Stack
#endif

#ifdef DEBUG
type Dbg = HasCallStack

debug :: Monad m => [String] -> m ()
debug strs = traceM (intercalate " | " strs ++ " END")

debugging :: Monad m => m () -> m ()
debugging act = act
{-# inline debugging #-}
#else
type Dbg = () :: Constraint

debugM :: Monad m => [String] -> m ()
debugM _ = pure ()
{-# inline debugM #-}

debug :: Monad m => [String] -> m ()
debug strs = pure ()
{-# inline debug #-}

debugging :: Monad m => m () -> m ()
debugging _ = pure ()
{-# inline debugging #-}
#endif

debug' :: Monad m => [String] -> m ()
debug' strs = traceM (intercalate " | " strs ++ " END")

debugging' :: Monad m => m () -> m ()
debugging' act = act
{-# inline debugging' #-}

-- unsafe performing
----------------------------------------------------------------------------------------------------

{-# noinline noinlineRunIO #-}
noinlineRunIO :: IO a -> a
noinlineRunIO (IO f) = runRW# (\s -> case f s of (# _, a #) -> a)

{-# inline runIO #-}
runIO :: IO a -> a
runIO (IO f) = runRW# (\s -> case f s of (# _, a #) -> a)

-- IO
----------------------------------------------------------------------------------------------------

type RW = State# RealWorld

{-# inline unIO #-}
unIO :: IO a -> RW -> (# RW, a #)
unIO = \(IO f) -> f

-- hole & error
----------------------------------------------------------------------------------------------------

{-# noinline uf #-}
uf :: Dbg => a
uf = undefined

{-# noinline impossible #-}
impossible :: Dbg => a
impossible = error "impossible"

-- Overloaded application
----------------------------------------------------------------------------------------------------

-- TODO: we don't need the separate lazy operator,
-- because Lazy is now in the type everywhere!
infixl 8 ∙
class Apply a b c constr | a b c -> constr where
  (∙) :: constr => a -> b -> c

-- strictness & primops
----------------------------------------------------------------------------------------------------

-- This is a diabolical way to override Prelude's ($) to be strict
{-# RULES "strictDollar" forall f x. f $ x = f $! x #-}

-- | If you want to go back to lazy '($)'.
{-# inline ($~) #-}
($~) :: (a -> b) -> a -> b
($~) = \f ~a -> f a

{-# inline ($$!) #-}
infixl 9 $$!
($$!) :: (a -> b) -> a -> b
($$!) f a = f $! a

-- for less annoying strict "idioms"
{-# inline (!) #-}
infixl 8 !
(!) :: Monad m => (a -> b) -> m a -> m b
(!) = (<$!>)

{-# inline (∘) #-}
infixl 8 ∘
(∘) :: Monad m => m (a -> b) -> m a -> m b
(∘) = \mf ma -> do f <- mf; a <- ma; pure $! f a

-- strict pairing
{-# inline (//) #-}
infixr 4 //
(//) :: a -> b -> (a, b)
a // b = (a, b)

{-# inline ptrEq #-}
ptrEq :: a -> a -> Bool
ptrEq x y = isTrue# (reallyUnsafePtrEquality# x y)

data Lazy a = Lazy {frc :: ~a} deriving Show

-- Not printing stuff by Show
----------------------------------------------------------------------------------------------------

newtype DontShow a = DontShow {unDontShow :: a} deriving (Eq, Ord)

instance Show (DontShow a) where
  showsPrec _ _ x = x

-- De Bruijn
----------------------------------------------------------------------------------------------------

newtype Ix = Ix {unIx :: Word}
  deriving (Eq, Ord, Show, Num, Enum, Bits, Integral, Real) via Word

newtype Lvl = Lvl {unLvl :: Word}
  deriving (Eq, Ord, Show, Num, Enum, Bits, Integral, Real) via Word

type LvlArg = (?lvl :: Lvl)

{-# inline setLvl #-}
setLvl :: Lvl -> (LvlArg => a) -> a
setLvl l k = let ?lvl = l in k

{-# inline lvlToIx #-}
lvlToIx :: LvlArg => Lvl -> Ix
lvlToIx (Lvl x) = Ix (coerce ?lvl - x - 1)

{-# inline ixToLvl #-}
ixToLvl :: LvlArg => Ix -> Lvl
ixToLvl (Ix x) = Lvl (coerce ?lvl - x - 1)


-- Source files
--------------------------------------------------------------------------------

data Src
  = SrcFile FilePath (RL.Ref ByteString) -- ^ Actual source file. It's under a Ref because
                                         --   reloading can mutably update the file.
  | SrcNoFile ByteString -- ^ Text coming from interactive input (CLI/LSP) or vanilla strings.

instance Show Src where
  showsPrec p = \case
    SrcFile path s -> showParen (p > 10) (("SrcFile " ++) . (path++))
    SrcNoFile s -> ("SrcInteractive"++)

srcToBs :: Src -> ByteString
srcToBs = \case
  SrcFile _ bs      -> runIO $ RL.read bs
  SrcNoFile bs -> bs

type SrcArg = (?src :: Src)

-- Source positions & spans
----------------------------------------------------------------------------------------------------

-- | A position is a mutable reference to a byte offset into an 'Src'. The 'Src' is stored somewhere
--   else in program context. The offset is counted *backwards* from the end of the source. Why it
--   is mutable: when re-parsing files after edits, we compare newly parsed presyntax to the old
--   presyntax. Most of the presyntax is usually unchanged, so we can skip a lot of elaboration.
--   However, positions might shift, so we update them while comparing preterms.
newtype Pos = Pos (RF.Ref FP.Pos)
  -- deriving Show via DontShow Pos

instance Show Pos where
  showsPrec p x =
    showParen (p>10) (("Pos "++).(case posToFPPos x of FP.Pos x -> (show x++)))

posToFPPos :: Pos -> FP.Pos
posToFPPos (Pos x) = runIO (RF.read x)

fpPosToPosIO :: FP.Pos -> IO Pos
fpPosToPosIO x = Pos ! RF.new x

-- | Equality of positions is stable under updates.
instance Eq Pos where
  Pos x == Pos y = runIO $ (==) ! RF.read x ∘ RF.read y

-- | Note: Ordering of positions is __not__ required to be stable under updates.
--   We only use this instance for error rendering.
instance Ord Pos where
  compare (Pos x) (Pos y) = runIO $ compare ! RF.read y ∘ RF.read y

deriving instance Flat FP.Pos

data Span = Span Pos Pos
  deriving Show via DontShow Span

fpSpanToSpanIO :: FP.Span -> IO Span
fpSpanToSpanIO (FP.Span x y) = Span ! fpPosToPosIO x ∘ fpPosToPosIO y

fpSpanToSpan :: FP.Span -> Span
fpSpanToSpan s = runIO (fpSpanToSpanIO s)

spanToRawName :: SrcArg => Span -> IO RawName
spanToRawName (Span (Pos i) (Pos j)) = do
  i <- RF.read i
  j <- RF.read j
  let bstr = srcToBs ?src
      i'   = B.length bstr - coerce i   -- FP.Pos counts backwards from the end of the string
      j'   = B.length bstr - coerce j
  pure $ RawName (B.take (j' - i') (B.drop i' bstr))

spanToString :: SrcArg => Span -> IO String
spanToString s = do
  RawName s <- spanToRawName s
  pure $ FP.utf8ToStr s

class SpanOf a where
  spanOf  :: a -> Span
  spanOf a = Span (leftPos a) (rightPos a)

  leftPos :: a -> Pos
  leftPos a = case spanOf a of Span l _ -> l

  rightPos :: a -> Pos
  rightPos a = case spanOf a of Span _ r -> r
  {-# minimal spanOf | leftPos, rightPos #-}

instance SpanOf Span where
  spanOf s = s

instance SpanOf Pos where
  leftPos  x = x
  rightPos x = x

instance SpanOf a => SpanOf (Maybe a) where
  leftPos = maybe impossible leftPos
  rightPos = maybe impossible rightPos

-- Parse errors & result types
----------------------------------------------------------------------------------------------------

data ParseErrorItem
  = PEIExpected String
  | PEILexicalError
  | PEIUnknown
  deriving Show

instance IsString ParseErrorItem where
  fromString = PEIExpected

data ParseError = ParseError [ParseErrorItem] Pos Pos Pos
  deriving Show

instance SpanOf ParseError where
  leftPos  (ParseError _ x _ _) = x
  rightPos (ParseError _ _ _ x) = x

data ParseMaybe a
  = PMJust a
  | PMNothing Pos
  deriving Show

class EmbedParseError a where
  embedParseError :: ParseMaybe a -> ParseError -> a

data ParsePair a b
  = PPair a b
  | PPError (ParseMaybe (ParsePair a b)) {-# nounpack #-} ParseError
  deriving Show

data ParseList a
  = PLNil Pos
  | PLCons a (ParseList a)
  | PLError (ParseMaybe (ParseList a)) ParseError
  deriving Show

pattern PLSingle :: SpanOf a => a -> ParseList a
pattern PLSingle a <- PLCons a (PLNil _) where
  PLSingle a = PLCons a (PLNil (rightPos a))
{-# inline PLSingle #-}

instance EmbedParseError (ParseList a) where
  embedParseError = PLError

instance EmbedParseError (ParsePair a b) where
  embedParseError = PPError

instance SpanOf a => SpanOf (ParseMaybe a) where
  leftPos  (PMJust a)    = leftPos a
  leftPos  (PMNothing p) = p
  rightPos (PMJust a)    = rightPos a
  rightPos (PMNothing p) = p

instance (SpanOf a, SpanOf b) => SpanOf (ParsePair a b) where
  leftPos  (PPair x _)   = leftPos x
  leftPos  (PPError x _) = leftPos x
  rightPos (PPair _ x)   = rightPos x
  rightPos (PPError _ x) = rightPos x

instance SpanOf a => SpanOf (ParseList a) where
  leftPos  (PLNil p)     = p
  leftPos  (PLCons x _)  = leftPos x
  leftPos  (PLError x _) = leftPos x
  rightPos (PLNil p)     = p
  rightPos (PLCons _ x)  = rightPos x
  rightPos (PLError _ x) = rightPos x

-- Names
----------------------------------------------------------------------------------------------------

-- | A 'ByteString' wrapper with potentially different instances.
newtype RawName = RawName ByteString
  deriving Eq via ByteString

instance Show RawName where
  show (RawName s) = FP.utf8ToStr s

-- TODO: use faster Hashable, Eq

instance Hashable RawName where
  hashWithSalt h (RawName (B.PS (ForeignPtr p _) (I# start) (I# len))) = go h (plusAddr# p start) len where
    goBytes :: Int -> Addr# -> Int# -> Int
    goBytes h p len = case len of
      0# -> h
      _  -> goBytes (hashWithSalt h (W8# (indexWord8OffAddr# p 0#))) (plusAddr# p 1#) (len -# 1#)

    go :: Int -> Addr# -> Int# -> Int
    go h p len = case len <# 8# of
      1# -> goBytes h p len
      _  -> go (hashWithSalt h (W64# (indexWord64OffAddr# p 0#))) (plusAddr# p 8#) (len -# 8#)


-- | Name coming from the user. Stores the source span plus the backing source data. Both are
--   mutable refs, because re-parsing a file can shift spans and it replaces the source
--   ByteString. NOTE: the utf8 data the span points to must never change.
data UserName = UserName Span Src

userNameToRawName :: UserName -> RawName
userNameToRawName (UserName s src) = let ?src = src in runIO $ spanToRawName s

instance Eq UserName where
  x == y = userNameToRawName x == userNameToRawName y

instance SpanOf UserName where
  spanOf (UserName s _) = s

instance Hashable UserName where
  hashWithSalt h x = hashWithSalt h $ userNameToRawName x

instance Show UserName where
  show x = show $ userNameToRawName x

data Name
  = NMachine ByteString -- ^ Name given by machine.
  | NUser UserName      -- ^ Name given by user.
  | N_                  -- ^ No name.
  deriving Eq

instance Hashable Name where
  hashWithSalt h = \case
    NMachine x -> (h + 1) `hashWithSalt` x
    NUser x    -> (h + 2) `hashWithSalt` x
    N_         -> h + 4

instance Show Name where
  showsPrec p (NMachine x) acc = showsPrec p x acc
  showsPrec p (NUser x)    acc = showsPrec p x acc
  showsPrec p N_           acc = '_':acc
  -- showsPrec p (NOp op)     acc = showsPrec p op acc

a_ = NMachine "a"
b_ = NMachine "b"
c_ = NMachine "c"
d_ = NMachine "d"
e_ = NMachine "e"
f_ = NMachine "f"
g_ = NMachine "g"
h_ = NMachine "h"
i_ = NMachine "i"
j_ = NMachine "j"
k_ = NMachine "k"
l_ = NMachine "l"
m_ = NMachine "m"
n_ = NMachine "n"
o_ = NMachine "o"
p_ = NMachine "p"
q_ = NMachine "q"
r_ = NMachine "r"
s_ = NMachine "s"
t_ = NMachine "t"
u_ = NMachine "u"
x_ = NMachine "x"
y_ = NMachine "y"
z_ = NMachine "z"

-- | Pick the more informative name.
pickName :: Name -> Name -> Name
pickName N_ N_ = x_
pickName N_ y  = y
pickName x N_  = x
pickName x y   = x

type Names = List Name
type NamesArg = (?names :: Names)

--------------------------------------------------------------------------------

class Sized a where
  size :: a -> Lvl

instance Sized (List a) where
  size = go 0 where
    go acc Nil = acc
    go acc (Cons _ as) = go (acc + 1) as

instance Sized [a] where
  size as = Lvl $ fromIntegral (length as)

type IsFlex = Bool

-- | Type for decision algorithm results that may get blocked by Flex things
data FlexDecide
  = Yes
  | No
  | Dunno
  deriving (Eq, Show)

--------------------------------------------------------------------------------

newtype MetaVar = MkMetaVar Int
  deriving (Eq, Ord, Num, Flat) via Int

instance Show MetaVar where
  showsPrec _ (MkMetaVar x) acc = '?': showsPrec 0 x acc

data PiDomKind
  = Expl   -- ^ Explicit checked argument
  | Impl   -- ^ Implicit argument
  | InfArg -- ^ Explicit inferred argument
  deriving (Eq, Show)

-- Overloaded fields
--------------------------------------------------------------------------------

class HasName s a | s -> a where
  name :: Lens' s a
  {-# minimal name #-}

class HasValue s a | s -> a where
  value :: Lens' s a
  {-# minimal value #-}

class HasBody s a | s -> a where
  body :: Lens' s a
  {-# minimal body #-}

class HasTy s a | s -> a where
  ty :: Lens' s a
  {-# minimal ty #-}