{-# LANGUAGE OverloadedStrings #-} module Main where import Database.MongoDB as DB import Text.HTML.TagSoup import System.Environment import System.IO import System.Process import Text.Printf import System.Exit import Control.Monad.Trans import Network.URI import Data.List.Split import Data.Maybe import Data.List (isPrefixOf,nub) import Data.Char import qualified Data.Text as T import qualified Data.Map as Map import Control.Monad import Data.Classify.DataTypes import qualified Data.Classify.Parser as Classify import qualified Data.Classify.Rank as Classify import Epguides main :: IO () main = do args <- getArgs case args of [addr, port, "magnet"] -> withMeteor addr port run [addr, port, "shows"] -> withMeteor addr port (getTVShows >>= liftIO . mapM_ putStrLn) [addr, port, "add", name] -> withMeteor addr port (addTVShow name) [addr, port, "remove", name] -> withMeteor addr port (removeTVShow name) [addr, port, "epguides", name] -> withMeteor addr port (mkEpguidesEntries name) [addr, port, "set_epguides", name, url] -> withMeteor addr port (setEpguides name url) [addr, port, "rebuild"] -> withMeteor addr port rebuild _ -> printUsage withMeteor addr port fn = do pipe <- runIOE $ connect (Host addr (PortNumber (fromIntegral (read port :: Integer)))) access pipe master "meteor" fn close pipe printUsage :: IO () printUsage = do prog <- getProgName hPutStrLn stderr $ printf "Usage: %s \n" prog hPutStrLn stderr $ printf "Commands:" hPutStrLn stderr $ printf " magnet Fetch magnet links from stdin and add to DB." hPutStrLn stderr $ printf " shows List the recognized TV shows." hPutStrLn stderr $ printf " add Add a name to the list of known TV shows." hPutStrLn stderr $ printf " remove Remove a name to the list of known TV shows." hPutStrLn stderr $ printf " epguides Fetch titles and air-times from stdin." hPutStrLn stderr $ printf " set_epguides Associate with ." hPutStrLn stderr $ printf " rebuild Recalculate computed properties in the DB." exitWith ExitSuccess -- stripName "The Big Bang Theory" = "bigbangtheory" -- stripName "CSI: Miami" = "csimiami" stripName :: String -> String stripName = worker . filter isAlpha . filter (not . isSpace) . map toLower where worker lst | "the" `isPrefixOf` lst = worker (drop 3 lst) worker lst = lst isMagnetLink :: URI -> Bool isMagnetLink uri = uriScheme uri == "magnet:" getDownloadName :: URI -> String getDownloadName uri = fromMaybe "" $ listToMaybe $ [ dn | ["dn", dn] <- map (splitOn "=") $ splitOn "&" (dropWhile (== '?') $ uriQuery uri) ] getHash :: URI -> String getHash uri = fromMaybe "" $ listToMaybe $ [ xt | ["xt", xt] <- map (splitOn "=") $ splitOn "&" (dropWhile (== '?') $ uriQuery uri) ] getTVShows = do records <- rest =<< find (select [] "meta") {project = [ "name" =: (1::Integer) ]} return $ nub $ mapMaybe (DB.lookup "name") records addTVShow name = insert "meta" ["name" =: name] setEpguides name url = modify (select ["name" =: name] "meta") ["$set" =: ["epguides" =: url]] getEpguidesURL name = do entry <- findOne (select ["name" =: name] "meta") case DB.lookup "epguides" =<< entry of Nothing -> error $ "No epguides URL associated with " ++ show name Just url -> return url removeTVShow name = delete (select ["name" =: name] "meta") getTags :: String -> [String] getTags = words . lower where lower str = [ if isAlphaNum c then c else ' ' | c <- str ] run = do inp <- liftIO $ getContents let hrefs = [ href | TagOpen "a" attrs <- parseTags inp, ("href", href) <- attrs ] shows <- getTVShows insertMagnetLinks shows hrefs insertMagnetLinks shows uris = do let magnets = [ uri | Just uri <- map parseURI uris, isMagnetLink uri ] idx = Map.fromList [ (stripName name, name) | name <- shows ] elts = [ (uri, Classify.selectBest elts) | uri <- magnets , let dn = getDownloadName uri , let elts = Classify.run (Classify.parseElement idx) dn , not (null elts) ] forM_ elts $ \(uri, element) -> do liftIO $ putStrLn $ "Adding element: " ++ show element let query = [ "name" =: name element , case version element of Version season episode -> "episode" =: [ "season" =: season, "episode" =: episode ] DateVersion year month day -> "airdate" =: [ "year" =: year, "month" =: month, "day" =: day ] ] inst = [ "link" =: show uri, "hash" =: getHash uri, "tags" =: getTags (title element) ] modify (select query "episodes") ["$addToSet" =: ["instances" =: inst]] return () rebuild = do shows <- getTVShows records <- rest =<< find (select ["instances" =: ["$exists" =: True]] "episodes") forM_ records $ \record -> do let record_id = DB.lookup "_id" record :: Maybe ObjectId forM_ (fromMaybe [] (DB.lookup "instances" record)) $ \inst -> case DB.lookup "link" inst of Nothing -> return () Just uri -> do modify (select ["_id" =: record_id] "episodes") ["$pull" =: ["instances" =: inst]] insertMagnetLinks shows [uri] {- {name: TV show ,searchable: string ,airdate: {year: int, month: int, day: int} ,episode: {season: int, episode: int} ,title: string ,instances: [{link: magnet, hash: hash, tags: [string]}] } -} mkEpguidesEntries name = do shows <- getTVShows unless (name `elem` shows) $ error $ "Unknown TV series: " ++ name epguides_url <- getEpguidesURL name inp <- liftIO $ readProcess "wget" ["-q", "-O", "-", epguides_url] [] forM_ (parseEpguide inp) $ \(EpIdx season episode, (year, month, day), title) -> do let query = ["name" =: name, "episode" =: ["season" =: season, "episode" =: episode ]] airdate = ["year" =: year, "month" =: month, "day" =: day] n <- count (select query "episodes") case n of 0 -> do liftIO $ putStrLn $ "Inserting new: " ++ show (season, episode, year, month, day, title) insert_ "episodes" [ "name" =: name , "searchable" =: map stripName (words name) , "episode" =: ["season" =: season, "episode" =: episode] , "airdate" =: airdate , "title" =: title ] _ -> do liftIO $ putStrLn $ "Updating: " ++ show (season, episode, year, month, day, title) modify (select query "episodes") ["$set" =: ["title" =: title, "airdate" =: airdate]]