module Data.Byte
    ( ByteOrder (..)
    , hostByteOrder
    , networkByteOrder
    , changeByteOrder
    , hton
    , ntoh

    , byteShow
    , byteRead
    , word8ToChar
    , charToWord8
    , showBits

    , logShift
    , logShiftL
    , logShiftR

    , flipEndian

    , word8Split
    , word8Concat
      
    , aton
    , ntoa )
    where


import Control.Exception

import Data.Bits
import Data.Char
import Data.List
import Data.Word

import Network.Socket

import System.Info
import System.IO.Unsafe

import Foreign.Marshal.Utils ( with )
import Foreign.Storable ( peekByteOff )

import Text.Regex


data ByteOrder = BigEndian | LittleEndian
                 deriving ( Eq, Show, Read )


hostByteOrder :: ByteOrder
hostByteOrder = let test = (word8Concat [1,2]) :: Word16
                    answer = (unsafePerformIO $ with test firstByte) :: Word8 in
                case answer of
                  1 -> LittleEndian
                  2 -> BigEndian
                  otherwise -> throw $ ErrorCall $ "Unexpected result when checking byte order"
    where firstByte = (flip peekByteOff) 0

networkByteOrder :: ByteOrder
networkByteOrder = BigEndian


changeByteOrder :: (Bits a, Integral a) => ByteOrder -> ByteOrder -> a -> a
changeByteOrder bo1 bo2 x = if bo1 == bo2
                            then x
                            else flipEndian x


-- |Host to Network byteorder
hton :: (Bits a, Integral a) => a -> a
hton = changeByteOrder hostByteOrder networkByteOrder 

-- |Network to Host byteorder
ntoh :: (Bits a, Integral a) => a -> a
ntoh = hton


-- |Converts the argument to a string representing the bytes in the 
-- argument
byteShow :: (Integral a, Bits a) => a -> String
byteShow = (map word8ToChar) . word8Split

-- |Reads a string of bytes. Reads only as many bytes as are needed
-- to represent the resulting type. Does not work with Integers
byteRead :: (Bits a, Integral a) => String -> a
byteRead = word8Concat . (map charToWord8)

word8ToChar :: Word8 -> Char
word8ToChar = chr . fromEnum

charToWord8 :: Char -> Word8
charToWord8 = toEnum . ord

-- |Returns a string representing the bits of the argument.
-- Does not work with Integers
showBits :: Bits a => a -> String
showBits x = showBits' ((bitSize x) -1) x
    where showBits' i x | i < 0 = ""
                        | testBit x i = '1' : showBits' (i-1) x
                        | otherwise = '0' : showBits' (i-1) x


-- |Reverses the endianness of the argument. Does not work
-- with Integers
flipEndian :: Bits a => a -> a
flipEndian x = bwOr $ zipWith logShift bytes shiftLengths
    where bytes = map (x .&.) $ map (shift 255) [0,8..sz-8]
          shiftLengths = [sz-8,sz-24..(negate sz)+8]
          sz = bitSize x

-- flipEndian = word8Concat . reverse . word8Split

bwOr :: Bits a => [a] -> a
bwOr = foldr (.|.) 0


-- |Logical shift version of the shift operator. Does not work with
-- Integers
logShift :: Bits a => a -> Int -> a
logShift x i | i >= 0 = logShiftL x i
             | otherwise = logShiftR x i

-- |Logical shift version of the shiftR operator. Does not work with
-- Integers
logShiftR :: Bits a => a -> Int -> a
logShiftR x i | isSigned x = clearBits (shiftR x i) (sz+i) sz
              | otherwise = shiftR x i
    where sz = bitSize x

-- |Identical to the shiftL operator.
logShiftL :: Bits a => a -> Int -> a
logShiftL = shiftL


clearBits :: Bits a => a -> Int -> Int -> a
clearBits x min max | min >= max = x
                    | otherwise = clearBits (clearBit x min) (min+1) max


-- |Returns a list of Word8s in little-endian byte order. Does not
-- work with Integers
word8Split :: (Integral a, Bits a) => a -> [Word8]
word8Split x = map (fromIntegral . (shiftR x)) [0,8..sz-8]
    where sz = bitSize x

-- |Concats a list of Word8s in little-endian byte order.
-- For big-endian byte order reverse the Word8s first.
-- Does not work with Integers
word8Concat :: (Bits a, Integral a) => [Word8] -> a
word8Concat w8s = bwOr $ zipWith shiftL xs [0,8..sz-8]
    where xs = map fromIntegral w8s
          sz = bitSize $ head xs


-- |Implementation of inet_addr outside the IO monad
aton :: String -> HostAddress
aton ipstr = word8Concat $ toBytes ipstr
    where toBytes = (map readByte) . checkLength . dotSplit
          dotSplit = splitRegex $ mkRegex "\\."
          checkLength x = if length x == 4 then x else malformed
          readByte b = case reads b of
                         ((x,""):_) -> if x >= 0 && x <= 255 
                                       then fromInteger x 
                                       else malformed
                         otherwise -> malformed
          malformed = throw $ ErrorCall $ "Malformed Address: " ++ ipstr


-- |Implementation of inet_ntoa outside the IO monad
ntoa :: HostAddress -> String
ntoa = concat . (intersperse ".") . (map show) . word8Split