import Data.Maybe (isJust)
import Data.List (nub)
import Control.Monad.Instances
import Debug.Trace (trace)

data Plan a = Plan a [Strategy a]
type Strategy a = Plan a -> Plan a

execute :: Plan a -> a
execute (Plan x []) = x
execute p@(Plan _ (s:_)) = execute (s p)

modify :: (a -> a) -> Strategy a
modify f (Plan x ss) = Plan (f x) ss

adjust :: ([Strategy a] -> [Strategy a]) -> Strategy a
adjust f (Plan x ss) = Plan x (f ss)

instead :: [Strategy a] -> Strategy a
instead ss = adjust (const ss)

done :: Strategy a
done = instead []

continue :: Strategy a
continue = adjust tail

result :: a -> Strategy a
result x = const (Plan x [])

resultWith :: (a -> a) -> Strategy a
resultWith f (Plan x ss) = Plan (f x) []

iff :: (a -> Bool) -> Strategy a -> Strategy a -> Strategy a
iff pred s1 s2 = \p@(Plan state fs) -> if (pred state) then (s1 p) else (s2 p)

state :: Plan a -> a
state (Plan a _) = a

computation :: Plan a -> [Strategy a]
computation (Plan _ ss) = ss

replace :: (a -> b) -> (b -> a -> a) -> Strategy a
replace get set (Plan s ss) = Plan s' ss
    where s' = set (get s) s

data Piece = X | O deriving (Eq,Show)
type Square = Maybe Piece
type Board = [Square]
type Line = [Int]

wins :: [Line]
wins = [[0,1,2],[3,4,5],[6,7,8],[0,3,6],[1,4,7],[2,5,8],[0,4,8],[2,4,6]]

get :: Board -> Int -> Square
get = (!!)

has2 :: Line -> Board -> Maybe Int
has2 line b = if (exactly2full && allSameColor) then empty else Nothing
    where exactly2full = length line' == 2
          allSameColor = (==1) . length . nub $ line'
          line' = filter isJust (map (get b) line)
          empty = Just . head . filter (not . isJust . get b) $ line

data TTTState = TTTState {
      board :: Board,
      turn :: Piece,
      move :: Maybe Int
} deriving Show

rememberMove :: Maybe Int -> TTTState -> TTTState
rememberMove m ttts = ttts { move = m }

findWinner :: Line -> Strategy TTTState
findWinner line = trace msg $ iff (canWin line) saveAndEnd continue
    where getWinner = has2 line . board
          canWin line = isJust . getWinner
          saveAndEnd = trace "saving..." $ resultWith $ rememberMove =<< getWinner
          msg = "checking " ++ show line

example :: TTTState
example = TTTState empty X Nothing
    where empty = replicate 9 Nothing

example2 = TTTState start X Nothing
    where start = (replicate 7 Nothing) ++ [Just X, Just X]

main = print $ execute (Plan example2 (map findWinner wins))

checking [0,1,2]
checking [3,4,5]
checking [6,7,8]
saving...
TTTState {board = [Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Just X,Just X], turn = X, move = Just 6}