{-# LANGUAGE BangPatterns, OverloadedStrings #-}
import Bio.Base
import Bio.Bam
import Bio.Bam.Evan ( removeWarts )
import Bio.Iteratee.ZLib
import Control.Monad
import Data.Bits
import System.Console.GetOpt
import System.Environment
import System.Exit
import System.IO

import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as S
import qualified Data.Vector.Generic as V

-- TODO:
-- - optional(!) GZip

data Opts = Opts { output :: BamMeta -> Iteratee [BamRec] IO ()
                 , inputs :: [Input]
                 , verbose :: Bool }

defaultOpts :: Opts
defaultOpts = Opts { output = protectTerm . pipeBamOutput
                   , inputs = []
                   , verbose = False }

data Input = Input { _read1 :: FilePath
                   ,  read2 :: Maybe FilePath
                   , index1 :: Maybe FilePath
                   , index2 :: Maybe FilePath }
  deriving Show

getopts :: [String] -> ([Opts -> IO Opts], [String], [String])
getopts = getOpt (ReturnInOrder add_read1) options
  where
    options =
        [ Option "o" ["output"] (ReqArg set_output "FILE") "Write output to FILE"
        , Option "1" ["read-one"] (ReqArg add_read1 "FILE") "Parse FILE for anything"
        , Option "2" ["read-two"] (ReqArg add_read2 "FILE") "Parse FILE for second mates"
        , Option "I" ["index-one"] (ReqArg add_idx1 "FILE") "Parse FILE for first index"
        , Option "J" ["index-two"] (ReqArg add_idx2 "FILE") "Parse FILE for second index"
        , Option "v" ["verbose"] (NoArg set_verbose) "Print progress information"
        , Option "h?" ["help","usage"] (NoArg usage) "Print this helpful message" ]

    set_output "-" c = return $ c { output = pipeBamOutput }
    set_output  fn c = return $ c { output = writeBamFile fn }
    set_verbose    c = return $ c { verbose = True }

    add_read1 fn c = return $ c { inputs = Input fn Nothing Nothing Nothing : inputs c }
    add_read2 fn c = return $ c { inputs = at_head (\i -> i { read2  = Just fn }) (inputs c) }
    add_idx1  fn c = return $ c { inputs = at_head (\i -> i { index1 = Just fn }) (inputs c) }
    add_idx2  fn c = return $ c { inputs = at_head (\i -> i { index2 = Just fn }) (inputs c) }

    at_head f [    ] = [ f $ Input "-" Nothing Nothing Nothing ]
    at_head f (i:is) = f i : is

    usage _ = do pn <- getProgName
                 let t = "Usage: " ++ pn ++ " [OPTION...]\n" ++
                         "Reads multiple FastA or FastQ files and converts them to BAM.  See manpage for details."
                 hPutStrLn stderr $ usageInfo t options
                 exitSuccess


main :: IO ()
main = do (opts, [], errors) <- getopts `fmap` getArgs
          unless (null errors) $ mapM_ (hPutStrLn stderr) errors >> exitFailure
          conf <- foldl (>>=) (return defaultOpts) opts
          pgm <- addPG Nothing

          let eff_inputs = if null (inputs conf) then [ Input "-" Nothing Nothing Nothing ] else inputs conf
          hPrint stderr $ eff_inputs

          foldr ((>=>) . enum_input) run (reverse eff_inputs) $
                joinI $ progress (verbose conf) $
                joinI $ mapChunks concatDuals $
                ilift liftIO $ output conf (pgm mempty)


type UpToTwo a = (a, Maybe a)

one :: a -> UpToTwo a
one a = (a, Nothing)

two :: a -> a -> UpToTwo a
two a b = (a, Just b)

mapU2 :: (a -> b) -> UpToTwo a -> UpToTwo b
mapU2 f (a,b) = (f a, fmap f b)

concatDuals :: [UpToTwo a] -> [a]
concatDuals ((a,Just  b):ds) = a : b : concatDuals ds
concatDuals ((a,Nothing):ds) = a : concatDuals ds
concatDuals [              ] = []

-- Enumerates a file.  Sequence and quality end up in b_seq and b_qual.
fromFastq :: (MonadIO m, MonadMask m) => FilePath -> Enumerator [BamRec] m a
fromFastq fp = enumAny fp $= enumInflateAny $= parseFastqCassava $= mapStream removeWarts
  where
    enumAny "-" = enumHandle defaultBufSize stdin
    enumAny  f  = enumFile defaultBufSize f

enum_input :: (MonadIO m, MonadMask m) => Input -> Enumerator [UpToTwo BamRec] m a
enum_input inp@(Input r1 mr2 mi1 mi2) o = do
    liftIO $ hPrint stderr inp
    (withIndex mi1 "XI" "YI" $ withIndex mi2 "XJ" "YJ" $
        case mr2 of Nothing -> fromFastq r1 $= mapStream one ; Just r2 -> enumDual r1 r2) o

-- Given an enumerator and maybe a filename, read index sequences from
-- the file and merge them with the numerator.
withIndex :: (MonadIO m, MonadMask m)
          => Maybe FilePath -> BamKey -> BamKey
          -> Enumerator [UpToTwo BamRec] m a -> Enumerator [UpToTwo BamRec] m a
withIndex Nothing      _    _ enum = enum
withIndex (Just fp) tagi tagq enum = mergeEnums enum (fromFastq fp) (convStream combine)
  where
    combine = do seqrecs <- lift headStream
                 idxrec  <- headStream
                 when (b_qname (fst seqrecs) /= b_qname idxrec) . error $
                        "read names do not match: " ++ shows (b_qname (fst seqrecs)) " & " ++ show (b_qname idxrec)

                 let idxseq  = S.pack $ map showNucleotides $ V.toList $ b_seq idxrec
                     idxqual = B.pack $ map   ((+33) . unQ) $ V.toList $ b_qual idxrec
                 return [ flip mapU2 seqrecs $
                        \r -> r { b_exts = (if B.null idxqual then id else insertE tagq (Text idxqual))
                                         $ insertE tagi (Text idxseq) $ b_exts r } ]

-- Enumerate dual files.  We read two FastQ files and match them up.  We
-- must make sure the names match, and we will flag everything as
-- 1st/2nd mate, no matter if the syntactic warts were present in the
-- files themselves.
enumDual :: (MonadIO m, MonadMask m)
         => FilePath -> FilePath -> Enumerator [UpToTwo BamRec] m a
enumDual f1 f2 = mergeEnums (fromFastq f1 $= mapStream one) (fromFastq f2) (convStream combine)
  where
    combine = do (firstMate, Nothing) <- lift headStream
                 secondMate           <- headStream

                 when (b_qname firstMate /= b_qname secondMate) . error $
                        "read names do not match: " ++ shows (b_qname firstMate) " & " ++ show (b_qname secondMate)

                 let qc = (b_flag firstMate .|. b_flag secondMate) .&. flagFailsQC
                     addx k = maybe id (updateE k) $ maybe (lookup k (b_exts secondMate)) Just $ lookup k (b_exts firstMate)
                     add_indexes = addx "XI" . addx "XJ" . addx "YI" . addx "YJ"

                 return [ two (firstMate  { b_flag = qc .|.  flagFirstMate .|. flagPaired .|. b_flag firstMate .&. complement flagSecondMate
                                          , b_exts = add_indexes $ b_exts firstMate })
                              (secondMate { b_flag = qc .|. flagSecondMate .|. flagPaired .|. b_flag secondMate .&. complement flagFirstMate
                                          , b_exts = add_indexes $ b_exts secondMate }) ]


progress :: MonadIO m => Bool -> Enumeratee [UpToTwo BamRec] [UpToTwo BamRec] m b
progress False = mapChunks id
progress True  = eneeCheckIfDonePass (icont . go 0 0)
  where
    go !_ !_ k (EOF         mx) = idone (liftI k) (EOF mx)
    go !l !n k (Chunk    [   ]) = liftI $ go l n k
    go !l !n k (Chunk as@(a:_)) = do
        let !n' = n + length as
            !nm = b_qname (fst a)
            !l' = l `max` S.length nm
        when (n `div` 2048 /= n' `div` 2048) $ liftIO $ do
            hPutStr stderr $ "\27[K" ++
                replicate (l' - S.length nm) ' '
                ++ S.unpack nm ++ ", "
                ++ shows n' " records processed\n"
            hFlush stderr
        eneeCheckIfDonePass (icont . go l' n') . k $ Chunk as