{-# LANGUAGE DoAndIfThenElse #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
module Main
where

-- System imports.
import Control.Arrow ( (&&&), (>>^), arr, returnA )
import Control.Concurrent ( threadDelay )
import Control.Exception ( SomeException, catch )
import Control.Monad ( when )
import Database.Groundhog.Generic ( runDbConn )
import Database.Groundhog.Sqlite (
  withSqliteConn )
import Database.Groundhog.Postgresql (
  withPostgresqlConn )
import Data.Monoid ( (<>) )
import Network.Services.TSN.Logging ( init_logging )
import System.Console.CmdArgs ( def )
import System.Directory ( removeFile )
import System.Exit ( exitWith, ExitCode (ExitFailure) )
import System.IO.Error ( catchIOError )
import Text.XML.HXT.Core (
  ArrowXml,
  IOStateArrow,
  SysConfigList,
  XmlTree,
  (>>>),
  (/>),
  getAttrl,
  getText,
  hasName,
  readDocument,
  runX,
  unpickleDoc )

-- Local imports.
import Backend ( Backend(..) )
import CommandLine ( get_args )
import Configuration ( Configuration(..), merge_optional )
import ConnectionString ( ConnectionString(..) )
import ExitCodes ( exit_no_xml_files )
import qualified OptionalConfiguration as OC (
  OptionalConfiguration ( xml_files ),
  from_rc )
import Network.Services.TSN.Report (
  report_info,
  report_error )
import TSN.DbImport ( DbImport(..), ImportResult(..) )
import TSN.Parse ( format_parse_error )
import qualified TSN.XML.AutoRacingDriverList as AutoRacingDriverList (
  dtd,
  pickle_message )
import qualified TSN.XML.AutoRacingResults as AutoRacingResults (
  dtd,
  pickle_message )
import qualified TSN.XML.AutoRacingSchedule as AutoRacingSchedule (
  dtd,
  pickle_message )
import qualified TSN.XML.EarlyLine as EarlyLine (
  dtd,
  pickle_message )
import qualified TSN.XML.GameInfo as GameInfo ( dtds, parse_xml )
import qualified TSN.XML.Heartbeat as Heartbeat ( dtd, verify )
import qualified TSN.XML.Injuries as Injuries ( dtd, pickle_message )
import qualified TSN.XML.InjuriesDetail as InjuriesDetail (
  dtd,
  pickle_message )
import qualified TSN.XML.MLBEarlyLine as MLBEarlyLine (
  dtd,
  pickle_message )
import qualified TSN.XML.JFile as JFile ( dtd, pickle_message )
import qualified TSN.XML.News as News (
  dtd,
  has_only_single_sms,
  pickle_message )
import qualified TSN.XML.Odds as Odds ( dtd, pickle_message )
import qualified TSN.XML.ScheduleChanges as ScheduleChanges (
  dtd,
  pickle_message )
import qualified TSN.XML.Scores as Scores ( dtd, pickle_message )
import qualified TSN.XML.SportInfo as SportInfo ( dtds, parse_xml )
import qualified TSN.XML.Weather as Weather (
  dtd,
  is_type1,
  pickle_message,
  teams_are_normal )
import Xml ( DtdName(..), parse_opts, parse_opts_novalidate )


-- | This is where most of the work happens. This function is called
--   on every file that we would like to import. It determines which
--   importer to use based on the DTD, attempts to process the file,
--   and then returns whether or not it was successful. If the file
--   was processed, 'True' is returned. Otherwise, 'False' is
--   returned.
--
--   The implementation is straightforward with one exception: since
--   we are already in arrow world with HXT, the @import_with_dtd@
--   function is lifted to an 'Arrow' as well with 'arr'. This
--   prevents us from having to do a bunch of unwrapping and
--   rewrapping with the associated error checking.
--
import_file :: Configuration -- ^ A configuration object needed for the
                             --   'backend' and 'connection_string'.

            -> FilePath -- ^ The path of the XML file to import.

            -> IO Bool -- ^ True if we processed the file, False otherwise.
import_file cfg path = do
  results <- parse_and_import `catch` exception_handler
  case results of
    []    -> do
      -- One of the arrows returned "nothing."  Now that we're
      -- validating against the DTDs, this will almost always be
      -- caused by a document whose DTD is not present (i.e. is
      -- unsupported). So we return "success" to allow the XML file to
      -- be deleted.
      report_info $ "No DTD for file " ++ path ++ "."
      return True
    (ImportFailed errmsg:_) -> do
      report_error $ errmsg ++ " (" ++ path ++ ")"
      return False
    (ImportSkipped infomsg:_) -> do
      -- We processed the message but didn't import anything. Return
      -- "success" so that the XML file is deleted.
      report_info infomsg
      return True
    (ImportSucceeded:_) -> do
      report_info $ "Successfully imported " ++ path ++ "."
      return True
    (ImportUnsupported infomsg:_) -> do
      -- For now we return "success" for these too, since we know we don't
      -- support a bunch of DTDs and we want them to get deleted.
      report_info infomsg
      return True
  where
    -- | This will catch *any* exception, even the ones thrown by
    --   Haskell's 'error' (which should never occur under normal
    --   circumstances).
    exception_handler :: SomeException -> IO [ImportResult]
    exception_handler e = do
      report_error (show e)
      let errdesc = "Failed to import file " ++ path ++ "."
      -- Return a nonempty list so we don't claim incorrectly that
      -- we couldn't parse the DTD.
      return [ImportFailed errdesc]

    -- | An arrow that reads a document into an 'XmlTree'.  We take a
    --   SysConfigList so our caller can decide whether or not to
    --   e.g. validate the document against its DTD.
    readA :: SysConfigList -> IOStateArrow s a XmlTree
    readA scl = readDocument scl path

    -- | An arrow which parses the doctype "SYSTEM" of an 'XmlTree'.
    --   We use these to determine the parser to use.
    dtdnameA :: ArrowXml a => a XmlTree DtdName
    dtdnameA = getAttrl >>> hasName "doctype-SYSTEM" /> getText >>^ DtdName

    -- | Combine the arrows above as well as the function below
    --   (arrowized with 'arr') into an IO action that does everything
    --   (parses and then runs the import on what was parsed).
    --
    --   The result of runX has type IO [IO ImportResult]. We thus use
    --   bind (>>=) and sequence to combine all of the IOs into one
    --   big one outside of the list.
    --
    --   Before we actually run the import, we check it against a list
    --   of problem DTDs. These can produce weird errors, and we have
    --   checks for them. But with DTD validation enabled, we can't
    --   even look inside the document to see what's wrong -- parsing
    --   will fail! So for those special document types, we proceed
    --   using 'parse_opts_novalidate' instead of the default
    --   'parse_opts'.
    --
    parse_and_import :: IO [ImportResult]
    parse_and_import = do
      -- Get the DTD name without validating against it.
      ((DtdName dtd) : _) <- runX $ (readA parse_opts_novalidate) >>> dtdnameA

      let problem_dtds = [ News.dtd, Weather.dtd ]
      let opts = if dtd `elem` problem_dtds
                 then parse_opts_novalidate
                 else parse_opts

      runX ((readA opts) >>> (dtdnameA &&& returnA) >>> (arr import_with_dtd))
        >>= sequence

    -- | Takes a ('DtdName', 'XmlTree') pair and uses the 'DtdName'
    --   to determine which function to call on the 'XmlTree'.
    import_with_dtd :: (DtdName, XmlTree) -> IO ImportResult
    import_with_dtd (DtdName dtd,xml)
        -- We special-case the heartbeat so it doesn't have to run in
        -- the database monad.
      | dtd == Heartbeat.dtd = Heartbeat.verify xml
      | otherwise =
        -- We need NoMonomorphismRestriction here.
        if backend cfg == Postgres
        then withPostgresqlConn cs $ runDbConn importer
        else withSqliteConn cs $ runDbConn importer
        where
          -- | Pull the real connection String out  of the configuration.
          --
          cs :: String
          cs = get_connection_string $ connection_string cfg

          -- | Convenience; we use this everywhere below in 'importer'.
          --
          migrate_and_import m = dbmigrate m >> dbimport m

          -- | The error message we return if unpickling fails.
          --
          errmsg = "Could not unpickle " ++ dtd ++ "."

          -- | Try to migrate and import using the given pickler @f@;
          --   if it works, return the result. Otherwise, return an
          --   'ImportFailed' along with our error message.
          --
          go f = maybe
                   (return $ ImportFailed errmsg)
                   migrate_and_import
                   (unpickleDoc f xml)

          importer
            | dtd == AutoRacingDriverList.dtd =
                go AutoRacingDriverList.pickle_message

            | dtd == AutoRacingResults.dtd =
                go AutoRacingResults.pickle_message

            | dtd == AutoRacingSchedule.dtd =
                go AutoRacingSchedule.pickle_message

            | dtd == EarlyLine.dtd =
                go EarlyLine.pickle_message

            -- GameInfo and SportInfo appear last in the guards
            | dtd == Injuries.dtd = go Injuries.pickle_message

            | dtd == InjuriesDetail.dtd = go InjuriesDetail.pickle_message

            | dtd == JFile.dtd = go JFile.pickle_message

            | dtd == MLBEarlyLine.dtd = go MLBEarlyLine.pickle_message

            | dtd == News.dtd =
                -- Some of the newsxml docs are busted in predictable ways.
                -- We want them to "succeed" so that they're deleted.
                -- We already know we can't parse them.
                if News.has_only_single_sms xml
                then go News.pickle_message
                else do
                  let msg = "Unsupported newsxml.dtd with multiple SMS " ++
                            "(" ++ path ++ ")"
                  return $ ImportUnsupported msg
            | dtd == Odds.dtd = go Odds.pickle_message

            | dtd == ScheduleChanges.dtd = go ScheduleChanges.pickle_message

            | dtd == Scores.dtd = go Scores.pickle_message

            -- SportInfo and GameInfo appear last in the guards
            | dtd == Weather.dtd =
                -- Some of the weatherxml docs are busted in predictable ways.
                -- We want them to "succeed" so that they're deleted.
                -- We already know we can't parse them.
                if Weather.is_type1 xml
                then if Weather.teams_are_normal xml
                     then go Weather.pickle_message
                     else do
                       let msg = "Teams in reverse order in weatherxml.dtd" ++
                                 " (" ++ path ++ ")"
                       return $ ImportUnsupported msg
                else do
                  let msg = "Unsupported weatherxml.dtd type (" ++ path ++ ")"
                  return $ ImportUnsupported msg

            | dtd `elem` GameInfo.dtds = do
                let either_m = GameInfo.parse_xml dtd xml
                case either_m of
                  -- This might give us a slightly better error
                  -- message than the default 'errmsg'.
                  Left err -> return $ ImportFailed (format_parse_error err)
                  Right m     -> migrate_and_import m

            | dtd `elem` SportInfo.dtds = do
                let either_m = SportInfo.parse_xml dtd xml
                case either_m of
                  -- This might give us a slightly better error
                  -- message than the default 'errmsg'.
                  Left err -> return $ ImportFailed (format_parse_error err)
                  Right m     -> migrate_and_import m

            | otherwise = do
              let infomsg =
                    "Unrecognized DTD in " ++ path ++ ": " ++ dtd ++ "."
              -- This should be an impossible case while DTD
              -- validation is enabled. If we can parse the file at
              -- all, then we have a DTD for it sitting around. And we
              -- only have DTDs for supported types.
              return $ ImportUnsupported infomsg



-- | Entry point of the program. It twiddles some knobs for
--   configuration options and then calls 'import_file' on each XML
--   file given on the command-line.
--
--   Any file successfully processed is then optionally removed, and
--   we're done.
--
main :: IO ()
main = do
  rc_cfg <- OC.from_rc
  cmd_cfg <- get_args

  -- Merge the config file options with the command-line ones,
  -- prefering the command-line ones.
  let opt_config = rc_cfg <> cmd_cfg

  -- Update a default config with any options that have been set in
  -- either the config file or on the command-line.  We initialize
  -- logging before the missing parameter checks below so that we can
  -- log the errors.
  let cfg = (def :: Configuration) `merge_optional` opt_config
  init_logging (log_level cfg) (log_file cfg) (syslog cfg)

  -- Check the optional config for missing required options.
  when (null $ OC.xml_files opt_config) $ do
    report_error "No XML files given."
    exitWith (ExitFailure exit_no_xml_files)

  -- We don't do this in parallel (for now?) to keep the error
  -- messages nice and linear.
  results <- mapM (import_file cfg) (OC.xml_files opt_config)

  -- Zip the results with the files list to find out which ones can be
  -- deleted.
  let result_pairs = zip (OC.xml_files opt_config) results
  let victims = [ p | (p, True) <- result_pairs ]
  let processed_count = length victims
  report_info $ "Processed " ++ (show processed_count) ++ " document(s) total."
  when (remove cfg) $ mapM_ (kill True) victims

  where
    -- | Wrap these two actions into one function so that we don't
    --   report that the file was removed if the exception handler is
    --   run.
    remove_and_report path = do
      removeFile path
      report_info $ "Removed processed file " ++ path ++ "."

    -- | Try to remove @path@ and potentially try again.
    kill try_again path =
      (remove_and_report path) `catchIOError` exception_handler
      where
        -- | A wrapper around threadDelay which takes seconds instead of
        --   microseconds as its argument.
        thread_sleep :: Int -> IO ()
        thread_sleep seconds = do
          let microseconds = seconds * (10 ^ (6 :: Int))
          threadDelay microseconds

        -- | If we can't remove the file, report that, and try once
        --   more after waiting a few seconds.
        exception_handler :: IOError -> IO ()
        exception_handler e = do
          report_error (show e)
          report_error $ "Failed to remove imported file " ++ path ++ "."
          if try_again then do
            report_info "Waiting 5 seconds to attempt removal again..."
            thread_sleep 5
            kill False path
          else
            report_info $ "Giving up on " ++ path ++ "."