{-# LANGUAGE MultiWayIf #-}
module Main where

import Control.Monad
import Data.Char
import Data.List
import Data.Ord
import qualified Data.Set as Set
import qualified Data.Map.Strict as Map
import System.IO
import Text.Parsec hiding (getInput)

import Debug.Trace

import Input
import qualified IntCode as IC


cyan, blue :: String -> String
cyan s = "\x1B[36m" ++ s ++ "\x1B[0m"
blue s = "\x1B[34m" ++ s ++ "\x1B[0m"


type Name = String
type Item = String

data Meta = Killed | Taken Item | WTF String
  deriving (Show)

data Dir = North | East | West | South
  deriving (Enum, Eq, Ord)

instance Show Dir where
    show North = "north"
    show East = "east"
    show West = "west"
    show South = "south"

instance Read Dir where
    readsPrec _ ('n':'o':'r':'t':'h':s) = [(North, s)]
    readsPrec _ ('e':'a':'s':'t':s) = [(East, s)]
    readsPrec _ ('w':'e':'s':'t':s) = [(West, s)]
    readsPrec _ ('s':'o':'u':'t':'h':s) = [(South, s)]
    readsPrec _ _ = []

data RoomDesc
    = RoomDesc { rName :: Name
               , rDesc :: String
               , rDoors :: [Dir]
               , rItems :: [Item] }
  deriving (Show)

parseOutput :: String -> Either Meta RoomDesc
parseOutput = either (\e -> Left (WTF (show e))) id . parse pRoom ""
  where
    pRoom = do
        spaces
        choice [do
                    void $ string "You take the "
                    item <- many (satisfy (/= '.'))
                    void $ char '.'
                    void $ newline
                    return (Left (Taken item))
               ,do
                    name <- pName
                    desc <- pDesc
                    spaces
                    void $ string "Doors here lead:\n"
                    dirStrs <- pList
                    spaces
                    items <- choice [string "Items here:\n" >> pList
                                    ,return []]
                    spaces
                    void $ string "Command?\n"
                    return (Right (RoomDesc name desc (map read dirStrs) items))
               ]
    pName = do
        void $ string "== "
        name <- many (notFollowedBy (try (string " ==")) >> anyChar)
        void $ string " =="
        void $ newline
        return name
    pDesc = pRestOfLine
    pList = many (string "- " >> pRestOfLine <* newline)
    pRestOfLine = many (satisfy (/= '\n'))

data Explorer = Explorer (Map.Map Name (RoomDesc, Map.Map Dir Name)) (Maybe Name) (Maybe (Name, Dir))
  deriving (Show)

observe :: Explorer -> RoomDesc -> Explorer
observe (Explorer chart _ lastMove) rd =
    let chart' = case Map.lookup (rName rd) chart of
                     Nothing -> Map.insert (rName rd) (rd, mempty) chart
                     Just _ -> chart
    in case lastMove of
           Nothing -> Explorer chart' (Just (rName rd)) lastMove
           Just (from, dir) ->
               let Just (fromRD, dirmap) = Map.lookup from chart
                   chart'' = Map.insert from (fromRD, Map.insert dir (rName rd) dirmap) chart'
               in Explorer chart'' (Just (rName rd)) lastMove

chooseMove :: Explorer -> Maybe (Dir, Explorer)
chooseMove (Explorer _ Nothing _) = undefined
chooseMove (Explorer chart (Just current@"Security Checkpoint") _) =
    Just (South, Explorer chart (Just current) (Just (current, South)))
chooseMove (Explorer chart (Just current) _) =
    let Just (rd, dirmap) = Map.lookup current chart
        opts = [case Map.lookup dir dirmap of
                    Just room' -> (dir, distToUnex room' 1 (Set.singleton current))
                    Nothing -> (dir, 1)
               | dir <- rDoors rd]
    in traceShow opts $
       case minimumBy (comparing snd) ((undefined, maxBound) : opts) of
           (_, n) | n == maxBound -> Nothing
           (dir, _) -> Just (dir, Explorer chart (Just current) (Just (current, dir)))
  where
    distToUnex :: Name -> Int -> Set.Set Name -> Int
    distToUnex "Security Checkpoint" _ _ = maxBound
    distToUnex room dist visited =
        let Just (rd, dirmap) = Map.lookup room chart
        in if | room `Set.member` visited ->
                  maxBound
              | not . null $ rDoors rd \\ Map.keys dirmap ->
                  dist + 1
              | otherwise ->
                  minimum [distToUnex room' (dist + 1) (Set.insert room visited)
                          | room' <- Map.elems dirmap]

printDOT :: Explorer -> String
printDOT (Explorer chart _ _) = "digraph G {\n" ++ printRelations "Hull Breach" mempty ++ "}"
  where
    printRelations room seen
        | room `Set.member` seen = ""
        | otherwise =
            let Just (rd, dirmap) = Map.lookup room chart
            in concat ((show room ++ " [label=" ++ show (room ++ "\n" ++ show (rItems rd)) ++ "];\n")
                      : unlines [show room ++ " -> " ++ show to ++ " [label=" ++ show (show dir) ++ "];"
                                | (dir, to) <- Map.assocs dirmap]
                      : [printRelations to (Set.insert room seen) | to <- Map.elems dirmap])

exploreAndPrintMap :: [Integer] -> IO ()
exploreAndPrintMap program =
    let loop cont inp expl =
            case IC.runContinue cont inp of
                Left (cont', out) -> do
                    let out' = map (chr . fromIntegral) out
                    putStr (cyan out') >> hFlush stdout
                    case parseOutput out' of
                        Right rd -> 
                            case chooseMove (observe expl rd) of
                                Just (dir, expl') -> do
                                    putStrLn (blue ("Entering '" ++ show dir ++ "'"))
                                    loop cont' (map (fromIntegral . ord) (show dir ++ "\n")) expl'
                                Nothing -> do
                                    putStrLn "Explorer gave up!"
                                    putStrLn (printDOT expl)
                        Left err -> do
                            print err
                            putStrLn (printDOT expl)
                Right (_, out) -> do
                    putStr (cyan (map (chr . fromIntegral) out))
    in loop (IC.initialContinuation program) [] (Explorer mempty Nothing Nothing)

interactive :: [Integer] -> IO ()
interactive program =
    let loop cont inp =
            case IC.runContinue cont inp of
                Left (cont', out) -> do
                    putStr (map (chr . fromIntegral) out)
                    line <- getLine
                    loop cont' (map (fromIntegral . ord) (line ++ "\n"))
                Right (_, out) -> do
                    putStr (map (chr . fromIntegral) out)
    in loop (IC.initialContinuation program) []

tryCombinations :: [Integer] -> IO ()
tryCombinations program =
    let loop _ _ [] = putStrLn "Out of tries!"
        loop cont inp (current:nexts) =
            case IC.runContinue cont inp of
                Left (cont', out) -> do
                    let out' = map (chr . fromIntegral) out
                    -- putStrLn out'
                    if | "lighter" `isInfixOf` out' -> putStrLn "  -> heavy"
                       | "heavier" `isInfixOf` out' -> putStrLn "  -> light"
                       | otherwise                  -> return ()
                    print current
                    let input = concat [concat ["take " ++ i ++ "\n" | i <- current]
                                       ,"east\n"
                                       ,concat ["drop " ++ i ++ "\n" | i <- current]]
                    loop cont' (map (fromIntegral . ord) input) nexts
                Right (_, out) -> do
                    putStr (map (chr . fromIntegral) out)
        setupString = "east\ntake antenna\nnorth\nnorth\ntake asterisk\nsouth\nwest\nwest\ntake astronaut ice cream\neast\nsouth\ntake hologram\nnorth\neast\nsouth\neast\ntake ornament\nnorth\nwest\ntake fixed point\neast\nsouth\nwest\nwest\nsouth\nsouth\nsouth\ntake dark matter\nnorth\nwest\nnorth\ntake monolith\nnorth\nnorth\ndrop astronaut ice cream\ndrop hologram\ndrop fixed point\ndrop asterisk\ndrop ornament\ndrop dark matter\ndrop monolith\ndrop antenna\n"
        allItems = ["asterisk", "ornament", "dark matter", "monolith", "antenna", "astronaut ice cream", "hologram", "fixed point"]
        combinations [] = [[]]
        combinations (x:xs) = let l = combinations xs in map (x:) l ++ l
    in loop (IC.initialContinuation program) (map (fromIntegral . ord) setupString) (combinations allItems)

main :: IO ()
main = do
    program <- IC.parse . head <$> getInput 25

    -- exploreAndPrintMap program
    -- interactive program
    tryCombinations program

    -- Collection of items: ["asterisk","ornament","astronaut ice cream","fixed point"]
    -- Password: 134227456