{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DerivingStrategies #-}

module Felix.Report.Location where

import Base
import Text.Megaparsec.Pos (SourcePos (sourceColumn, sourceLine), unPos)
import Data.Bits
import Control.DeepSeq (NFData)
import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef)
import Data.IntMap.Strict qualified as IntMap
import Data.Map.Strict qualified as Map
import Data.Text qualified as Text
import Data.Word (Word16, Word32)
import System.IO.Unsafe (unsafePerformIO)

-- | File identifier used in packed source locations.
newtype FileId = FileId
    { unFileId :: Word16
    } deriving stock (Show, Eq, Ord, Generic)
      deriving anyclass (Hashable, NFData)

-- | Packed source location.
-- Bit layout (high to low):
-- 16 bits file id | 24 bits line | 24 bits column.
newtype Location = Location
    { unLocation :: Word64
    } deriving stock (Eq, Ord, Generic)
      deriving anyclass (Hashable, NFData)

data RegisteredFile = RegisteredFile
    { registeredFileIdentity :: FilePath
    , registeredFileDisplayPath :: FilePath
    }

data FileRegistry = FileRegistry
    { registrationToFileId :: Map (FilePath, FilePath) FileId
    , fileIdToFile :: IntMap RegisteredFile
    , fileIdAllocator :: !FileIdAllocator
    }

data LocationRegistrationError
    = FileIdSpaceExhausted
    deriving stock (Show, Eq)

data LocationConstructionError
    = LocationCoordinateOutOfRange !Int !Int
    deriving stock (Show, Eq)

newtype FileIdAllocator = FileIdAllocator Word32
    deriving stock (Show, Eq)

initialFileIdAllocator :: FileIdAllocator
initialFileIdAllocator = FileIdAllocator 0

allocateFileId
    :: FileIdAllocator
    -> Either LocationRegistrationError (FileId, FileIdAllocator)
allocateFileId (FileIdAllocator next)
    | next >= fromIntegral (maxBound :: Word16) =
        Left FileIdSpaceExhausted
    | otherwise =
        Right
            ( FileId (fromIntegral next)
            , FileIdAllocator (next + 1)
            )

initialFileRegistry :: FileRegistry
initialFileRegistry = FileRegistry
    { registrationToFileId = mempty
    , fileIdToFile = mempty
    , fileIdAllocator = initialFileIdAllocator
    }

fileRegistryRef :: IORef FileRegistry
fileRegistryRef = unsafePerformIO (newIORef initialFileRegistry)
{-# NOINLINE fileRegistryRef #-}

registerFilePath
    :: MonadIO io
    => FilePath
    -> io (Either LocationRegistrationError FileId)
registerFilePath path =
    registerFilePathWithDisplay path path

-- | Register a file by physical identity while retaining a separate path for
-- diagnostics. The identity must be stable and unique for the selected
-- physical source. Distinct display paths receive distinct file identifiers so
-- one workspace cannot inherit another workspace's presentation. Callers with
-- only one path should use 'registerFilePath'.
registerFilePathWithDisplay
    :: MonadIO io
    => FilePath
    -> FilePath
    -> io (Either LocationRegistrationError FileId)
registerFilePathWithDisplay identity displayPath = liftIO $
    atomicModifyIORef' fileRegistryRef \registry ->
        case Map.lookup registration (registrationToFileId registry) of
            Just fileId ->
                (registry, Right fileId)
            Nothing ->
                case allocateFileId (fileIdAllocator registry) of
                    Left err ->
                        (registry, Left err)
                    Right (fileId, nextAllocator) ->
                        let fileIdInt = fromIntegral (unFileId fileId)
                            registeredFile =
                                RegisteredFile identity displayPath
                            registry' = FileRegistry
                                { registrationToFileId =
                                    Map.insert
                                        registration
                                        fileId
                                        (registrationToFileId registry)
                                , fileIdToFile =
                                    IntMap.insert
                                        fileIdInt
                                        registeredFile
                                        (fileIdToFile registry)
                                , fileIdAllocator = nextAllocator
                                }
                        in
                            (registry', Right fileId)
  where
    registration = (identity, displayPath)

lookupFilePath :: FileId -> Maybe FilePath
lookupFilePath fileId = unsafePerformIO do
    registry <- readIORef fileRegistryRef
    pure
        (registeredFileDisplayPath <$>
            IntMap.lookup
                (fromIntegral (unFileId fileId))
                (fileIdToFile registry))

lookupFileIdentityPath :: FileId -> Maybe FilePath
lookupFileIdentityPath fileId = unsafePerformIO do
    registry <- readIORef fileRegistryRef
    pure
        (registeredFileIdentity <$>
            IntMap.lookup
                (fromIntegral (unFileId fileId))
                (fileIdToFile registry))

fileShift, lineShift :: Int
fileShift = 48
lineShift = 24

fileMask, coordMask :: Word64
fileMask = 0xFFFF
coordMask = 0xFFFFFF

nowhereWord :: Word64
nowhereWord = maxBound

mkLocationChecked
    :: FileId
    -> Int
    -> Int
    -> Either LocationConstructionError Location
mkLocationChecked fileId line column
    | line < 0
        || column < 0
        || lineWord > coordMask
        || columnWord > coordMask =
            Left (LocationCoordinateOutOfRange line column)
    | otherwise =
        Right
            (Location
                (   (fromIntegral (unFileId fileId) `shiftL` fileShift)
                 .|. (lineWord `shiftL` lineShift)
                 .|. columnWord
                ))
    where
        lineWord = fromIntegral line :: Word64
        columnWord = fromIntegral column :: Word64

mkLocation :: FileId -> Int -> Int -> Location
mkLocation fileId line column =
    either
        (impossible . ("mkLocation: " <>) . show)
        id
        (mkLocationChecked fileId line column)

locFileId :: Location -> Maybe FileId
locFileId (Location w)
    | w == nowhereWord = Nothing
    | otherwise = Just (FileId (fromIntegral ((w `shiftR` fileShift) .&. fileMask)))

locFile :: Location -> FilePath
locFile loc = case locFileId loc of
    Nothing -> "<nowhere>"
    Just fileId ->
        fromMaybe ("<file#" <> show (unFileId fileId) <> ">") (lookupFilePath fileId)

locLine :: Location -> Int
locLine (Location w)
    | w == nowhereWord = -1
    | otherwise = fromIntegral ((w `shiftR` lineShift) .&. coordMask)

locColumn :: Location -> Int
locColumn (Location w)
    | w == nowhereWord = -1
    | otherwise = fromIntegral (w .&. coordMask)

instance Show Location where
    showsPrec p loc =
        showParen (p > appPrec) $
            showString "Location {locFile = "
            . shows (locFile loc)
            . showString ", locLine = "
            . shows (locLine loc)
            . showString ", locColumn = "
            . shows (locColumn loc)
            . showString "}"
      where
        appPrec = 10

fromSourcePosChecked
    :: FileId
    -> SourcePos
    -> Either LocationConstructionError Location
fromSourcePosChecked fileId pos =
    mkLocationChecked
        fileId
        (unPos (sourceLine pos))
        (unPos (sourceColumn pos))

prettyLocation :: Location -> String
prettyLocation loc =
    locFile loc <> " " <> show (locLine loc) <> ":" <> show (locColumn loc)

-- | Render two locations together, qualifying equal display paths when they
-- refer to different physical files.
prettyLocationPair :: Location -> Location -> (String, String)
prettyLocationPair firstLocation secondLocation =
    if locFile firstLocation == locFile secondLocation
        && firstIdentity /= secondIdentity
        && isJust firstIdentity
        && isJust secondIdentity
        then
            ( qualify firstIdentity firstLocation
            , qualify secondIdentity secondLocation
            )
        else
            (prettyLocation firstLocation, prettyLocation secondLocation)
  where
    firstIdentity =
        locFileId firstLocation >>= lookupFileIdentityPath
    secondIdentity =
        locFileId secondLocation >>= lookupFileIdentityPath

    qualify identity location =
        locFile location
            <> " (canonical "
            <> maybe "<unknown>" show identity
            <> ") "
            <> show (locLine location)
            <> ":"
            <> show (locColumn location)

locationToText :: Location -> Text
locationToText loc = Text.pack (prettyLocation loc)

-- | Things that have a location.
class Locatable a where
    locate :: a -> Location

pattern Nowhere :: Location
pattern Nowhere = Location 0xFFFFFFFFFFFFFFFF

instance Locatable Location where
    locate = id

instance Locatable a => Locatable [a] where
    locate []    = Nowhere
    locate (x:_) = locate x

instance Locatable a => Locatable (Maybe a) where
    locate Nothing  = Nowhere
    locate (Just x) = locate x

instance Locatable a => Locatable (NonEmpty a) where
    locate (x :| _) = locate x