module Main where

import Data.Char
import Data.List
import Control.Monad
import Data.Bits
import Data.Word
import qualified Data.Map as Map
import qualified Data.Set as Set

type Inttype = Word16

data Literal = LNum Inttype | LName String

data Command = Immediate Literal |
               Monadic (Inttype -> Inttype) Literal |
               Dyadic (Inttype -> Inttype -> Inttype) Literal Literal

type Statement = (Command,String)

type Cache = Map.Map String Inttype

complement16 :: Inttype -> Inttype
complement16 i = i `xor` 0xffff

shiftL16 :: Inttype -> Int -> Inttype
shiftL16 i n = (shiftL i n) .&. 0xffff

shiftR16 :: Inttype -> Int -> Inttype
shiftR16 i n = (shiftR i n) .&. 0xffff

parseLit :: String -> Literal
parseLit s = if isDigit (head s) then LNum (read s) else LName s

parse :: [String] -> Statement
parse l = (parse' (init $ init $ l),last l)

parse' :: [String] -> Command
parse' [n] = Immediate (parseLit n)
parse' ["NOT",n] = Monadic complement16 (parseLit n)
parse' [a,"AND",b] = Dyadic (.&.) (parseLit a) (parseLit b)
parse' [a,"OR",b] = Dyadic (.|.) (parseLit a) (parseLit b)
parse' [a,"LSHIFT",b] = Dyadic (\i n -> shiftL16 i (fromIntegral n)) (parseLit a) (parseLit b)
parse' [a,"RSHIFT",b] = Dyadic (\i n -> shiftR16 i (fromIntegral n)) (parseLit a) (parseLit b)

collectOutputs :: [Statement] -> [String]
collectOutputs stmts = Set.toList $ collectOutputs' stmts Set.empty

collectOutputs' :: [Statement] -> Set.Set String -> Set.Set String
collectOutputs' [] set = set
collectOutputs' ((_,target):ss) set = collectOutputs' ss $ Set.insert target set

neededBy :: String -> Command -> Bool
neededBy n (Immediate l) = case l of
    LNum _ -> False
    LName ln -> n == ln
neededBy n (Monadic _ l) = case l of
    LNum _ -> False
    LName ln -> n == ln
neededBy n (Dyadic _ (LNum _) b) = case b of
    LNum _ -> False
    LName bn -> n == bn
neededBy n (Dyadic _ (LName an) b) = case b of
    LNum _ -> n == an
    LName bn -> n == an || n == bn

leaves :: [Statement] -> [String]
leaves stmts = leaves' (map fst stmts) (collectOutputs stmts)

leaves' :: [Command] -> [String] -> [String]
leaves' _ [] = []
leaves' cmds (o:os) = if any (neededBy o) cmds then leaves' cmds os else o:leaves' cmds os

getValue :: String -> [Statement] -> Cache -> (Inttype,Cache)
getValue nm stmts cc = getValue' nm stmts stmts cc

getValue' :: String -> [Statement] -> [Statement] -> Cache -> (Inttype,Cache)
getValue' nm [] _ cc = case Map.lookup nm cc of
    Nothing -> undefined
    Just i -> (i,cc)
getValue' nm ((cmd,target):ss) stmts cc = case Map.lookup nm cc of
    Nothing -> if nm == target then execCmd cmd stmts cc else getValue' nm ss stmts cc
    Just i -> (i,cc)

resolveLit :: Literal -> [Statement] -> Cache -> (Inttype,Cache)
resolveLit (LNum i) _ cc = (i,cc)
resolveLit (LName n) stmts cc = (value,Map.insert n value cc2)
    where (value,cc2) = getValue n stmts cc

execCmd :: Command -> [Statement] -> Cache -> (Inttype,Cache)
execCmd (Immediate l) stmts cc = resolveLit l stmts cc
execCmd (Monadic f l) stmts cc = (f resl,cc2)
    where (resl,cc2) = resolveLit l stmts cc
execCmd (Dyadic f a b) stmts cc = (f resa resb,cc3)
    where (resa,cc2) = resolveLit a stmts cc
          (resb,cc3) = resolveLit b stmts cc2

day7 :: IO ()
day7 = do
    input <- readFile "input.txt"
    let stmts = [parse (words l) | l <- (lines input)]
    let lvs = leaves stmts
    print lvs
    print [fst $ getValue lv stmts Map.empty | lv <- lvs]

main = day7