{-# LANGUAGE TupleSections #-}

import Common
import AStar
import Data.List.Split (splitOn)
import Data.Maybe (catMaybes)
import Data.List (find, maximumBy)
import GHC.Data.Maybe (fromJust)
import GHC.Utils.Misc (sndOf3)

data Valve = Valve {
  valveName :: String,
  valveFlowRate :: Int,
  valveLeadsTo :: [String],
  valveOpen :: Bool
}
  deriving (Show, Eq)

-- disgusting, but works
parse :: String -> Valve
parse s = Valve name flowRate leadsTo False
  where
    leadsTo = map trim $ splitOn "," $ unwords $ drop 9 w
    name = w !! 1
    flowRate = read $ init $ (!! 1) $ splitOn "=" $ w !! 4
    w = words s

startLocation = "AA"
totalMinutes = 30

findValve :: [Valve] -> String -> Valve
findValve valves s = fromJust $ find ((== s) . valveName) valves

calculateValvesBuildup :: [Valve] -> Int
calculateValvesBuildup = sum . map valveFlowRate . filter valveOpen

openValve :: [Valve] -> String -> [Valve]
openValve valves name = map (\v -> if valveName v == name then Valve (valveName v) (valveFlowRate v) (valveLeadsTo v) True else v) valves

pathfindValve :: [Valve] -> String -> String -> (Int, [String])
pathfindValve valves a b = fromJust $ astarSearch a (== b) (map (, 1000) . valveLeadsTo . findValve valves) (const 0)

getValveScore :: [Valve] -> String -> Int -> Valve -> (Valve, Int, [String])
getValveScore valves pos minutes valve = (valve, (totalMinutes - approxMinutes) * flowRate, path)
  where
    flowRate = valveFlowRate valve
    approxMinutes = minutes + length path
    path = snd $ pathfindValve valves pos (valveName valve)

go :: [Valve] -> Int -> String -> Int -> Int
go valves minutes pos pressure
  | minutes > 30 = pressure
  | null valveScores = go valves (minutes + 1) pos (pressure + calculateValvesBuildup valves)
  | otherwise = maximum $ trace (map sndOf3 valveScores) $ map goTo maxValves
  where
    goTo :: (Valve, Int, [String]) -> Int
    goTo (valve, _, path) = go newValves (minutes + 1) newPosition (pressure + calculateValvesBuildup valves)
      where
        newPosition = if name == pos then pos else path !! 1
        newValves = if name == pos then openValve valves name else valves
        name = valveName valve

    maxValves = filter (\(_, score, _) -> score == maxScore) valveScores
    (_, maxScore, _) = maximumBy (\a b -> sndOf3 a `compare` sndOf3 b) valveScores
    valveScores = map (getValveScore valves pos minutes) $ filter (\v -> valveFlowRate v > 0 && not (valveOpen v)) valves

main = interact $
  show
  . (\valves -> go valves 0 startLocation 0)
  . map parse
  . lines