camh-0.0.2: camh.hs

{-# LANGUAGE OverloadedStrings #-}

import Control.Applicative ((<|>))
import Control.Monad
import Control.Exception (handle)
import Data.Bits
import qualified Data.ByteString.Char8 as B
import Data.List (group)
import Data.Word (Word32)
import Foreign.Ptr
import Foreign.Storable
import Graphics.Imlib
import System.Console.GetOpt
import System.Console.Terminfo
import System.Environment (getArgs)
import System.Exit (exitSuccess)
import System.IO (stderr, hPutStrLn)

from24bitColor :: Word32 -> Int
from24bitColor w =
    let p = fromIntegral w
        r = p `shift` (-16) .&. 255
        g = p `shift`  (-8) .&. 255
        b = p               .&. 255
    in 16 + 36 * (r `quot` 43) + 6 * (g `quot` 43) + (b `quot` 43)

changeBgColor256 :: Int -> B.ByteString
changeBgColor256 col =
    "\x1b[48;5;" `B.append` B.pack (show col) `B.append` "m"

restoreBgColor :: B.ByteString
restoreBgColor = "\x1b[m"

colorLine :: [Word32] -> B.ByteString
colorLine row =
    let g        = group row
        e        = changeBgColor256 . from24bitColor
        s n      = B.replicate n ' '
        f (w:ws) = e w `B.append` s (fromIntegral $ 1 + length ws)
    in B.concat (map f g) `B.append` restoreBgColor

writeImage :: Int -> Int -> Ptr Word32 -> IO ()
writeImage w h p =
    forM_ [0 .. h-1] $ \j -> do
      row <- forM [0 .. w-1] $ \i -> peekElemOff p (w * j + i)
      B.putStrLn $ colorLine row

dealFile :: Config -> FilePath -> IO ()
dealFile conf filename =
  loadImageWithErrorReturn filename >>= \(image, err) ->
      if err == ImlibLoadErrorNone then
          do
            let tw = getWidth  conf
                th = getHeight conf

            contextSetImage image
            iw <- imageGetWidth
            ih <- imageGetHeight

            let ratio x y = fromIntegral x / fromIntegral y
                z  = ratio tw iw `min` ratio (2 * th) ih
                rw = floor $ fromIntegral iw * z
                rh = floor $ fromIntegral ih * z / 2

            createCroppedScaledImage 0 0 iw ih rw rh >>= contextSetImage

            createColorModifier >>= contextSetColorModifier

            let whenJust (Just a) f = f a
                whenJust Nothing  f = return ()
            whenJust (getGamma      conf) modifyColorModifierGamma
            whenJust (getBrightness conf) modifyColorModifierBrightness
            whenJust (getContrast   conf) modifyColorModifierContrast

            applyColorModifier
            freeColorModifier

            imageWithData $ writeImage rw rh
            freeImage

            contextSetImage image
            freeImage
      else
          hPutStrLn stderr $ "error " ++ show err ++ ": " ++ filename

getTermSize :: IO (Maybe (Int, Int))
getTermSize = handle (const $ return Nothing
                          :: SetupTermError -> IO (Maybe a)) $ do
  term <- setupTermFromEnv
  let size = do
        lines <- tiGetNum "lines"
        cols  <- tiGetNum "cols"
        return (lines, cols)

  return $ getCapability term size

data Opts = Help
          | Version
          | WidthChr   Int
          | HeightChr  Int
          | Gamma      Double
          | Brightness Double
          | Contrast   Double

options :: Config -> [OptDescr Opts]
options conf =
    let tw = getWidth  conf
        th = getHeight conf
    in [ Option "?" ["help"] (NoArg Help) "Show this message"
       , Option "v" ["version"] (NoArg Version) "Show version info"
       , Option "W" ["width-char"]
                    (ReqArg (WidthChr . read) "WIDTH")
                    "Set output width (in chars)"
       , Option "w" ["width-percent"]
                    (ReqArg (WidthChr . percent tw . read) "WIDTH")
                    "Set output width (in percent)"
       , Option "H" ["height-char"]
                    (ReqArg (HeightChr . read) "HEIGHT")
                    "Set output height (in chars)"
       , Option "h" ["height-percent"]
                    (ReqArg (HeightChr . percent th . read) "HEIGHT")
                    "Set output height (in percent)"
       , Option "g" ["gamma"]
                    (ReqArg (Gamma . read) "GAMMA")
                    "Apply gamma correction"
       , Option "b" ["blightness"]
                    (ReqArg (Brightness . read) "BLIGHTNESS")
                    "Apply blightness correction"
       , Option "c" ["contrast"]
                    (ReqArg (Contrast . read) "CONTRAST")
                    "Apply contrast correction" ]

percent :: Int -> Double -> Int
percent v p = floor $ fromIntegral v * p / 100

usage = usageInfo (version ++ "\n\nOptions:") $ options defaultConfig

version = "camh version 0.0.2"

data Config = Config { getWidth      :: Int
                     , getHeight     :: Int
                     , getGamma      :: Maybe Double
                     , getBrightness :: Maybe Double
                     , getContrast   :: Maybe Double }

defaultConfig = Config { getWidth      = 80
                       , getHeight     = 25
                       , getGamma      = Nothing
                       , getBrightness = Nothing
                       , getContrast   = Nothing }

getOpts :: Config -> [String] -> IO ([Opts], [String])
getOpts conf args =
    let th = getHeight conf
        tw = getWidth  conf
    in case getOpt Permute (options conf) args of
         (o, n, []  ) -> return (o, n)
         (_, _, errs) -> error $ concat errs ++ usage

reduceOpt :: Config -> Opts -> IO Config
reduceOpt conf opts =
    case opts of
      Help         -> putStr   usage   >> exitSuccess
      Version      -> putStrLn version >> exitSuccess
      WidthChr   w -> return conf { getWidth      = w }
      HeightChr  h -> return conf { getHeight     = h }
      Gamma      g -> return conf { getGamma      = Just g }
      Brightness b -> return conf { getBrightness = Just b }
      Contrast   c -> return conf { getContrast   = Just c }

main = do
  s <- getTermSize
  let Just (th, tw) = s <|> Just ( getHeight defaultConfig
                                 , getWidth  defaultConfig )
      conf = defaultConfig { getWidth = tw, getHeight = th-1 }

  (opts, args) <- getArgs >>= getOpts conf

  conf' <- foldM reduceOpt conf opts

  mapM_ (dealFile conf') args