module Model where

import Data.Array.Unboxed
import Data.Map (Map)
import Data.Map qualified as Map
import Graphics.Vty ( Attr )
import Control.Monad ( guard, msum )
import Data.Maybe ( listToMaybe, mapMaybe, maybeToList )
import Data.Set qualified as Set
import Data.Set (Set)
import GHC.Stack (HasCallStack)

type Coord = (Int, Int)

data Box = Box {
  boxLocation :: Maybe Location,
  boxType     :: BoxType,
  boxColor    :: Attr,
  boxBoring   :: Bool
}
   deriving (Show, Read, Eq)

-- | Compute the height and width of a box.
boxSize :: World -> Box -> (Int, Int)
boxSize world box = (yhi-ylo+1, xhi-xlo+1)
  where
    ((ylo,xlo),(yhi,xhi)) = bounds (boxWalls world box)

-- | Find the array corresponding to the walls of a box.
boxWalls :: World -> Box -> UArray Coord Bool
boxWalls world box =
    case boxType box of
        Original walls  -> walls
        Link c          -> boxWalls world (boxIx world c)
        Infinity c      -> boxWalls world (boxIx world c)
        Epsilon _ walls -> walls

data BoxType
    = Original (UArray Coord Bool)      -- ^ plain box
    | Link     Char                     -- ^ entrance clone
    | Infinity Char                     -- ^ infinite exit
    | Epsilon  Char (UArray Coord Bool) -- ^ infinite enter
   deriving (Show, Read, Eq)

data Location = Location { locName :: Char, locY :: Int, locX :: Int }
  deriving (Read, Show, Ord, Eq)

data World = World {
  worldBoxes    :: Map Char Box,
  worldMe       :: Char,
  worldButtons  :: Set Location,
  worldHome     :: Location,
  worldSize     :: (Int, Int)
}
  deriving (Show, Read, Eq)

winCondition :: World -> Bool
winCondition world =
  Set.isSubsetOf (worldButtons world) coverage &&
  Just (worldHome world) == boxLocation (boxIx world (worldMe world))
  where
    coverage = Set.fromList (mapMaybe boxLocation (Map.elems (worldBoxes world)))

move :: World -> (Int,Int) -> World
move world dir =
  case myLocation world of
    Nothing -> world
    Just loc ->
      let locMap = worldLocations world in
      case moveBlock world locMap Map.empty loc dir 0 of
        Nothing -> world
        Just changes ->
          let f box change = box { boxLocation = change } in
          world { worldBoxes = Map.mergeWithKey (\_ a b -> Just (f a b)) id (const Map.empty) (worldBoxes world) changes}

myLocation :: World -> Maybe Location
myLocation world = boxLocation (boxIx world (worldMe world))

isWall :: World -> Location -> Bool
isWall world (Location n y x) =
  case Map.lookup n (worldBoxes world) of
    Nothing -> True
    Just box -> boxWalls world box ! (y,x)

type Movement = (Int, Int)

moveBlock ::
  World ->
  Map Location (Char, Box) ->
  Map Location (Int, Char, Maybe Location) ->
  Location ->
  Movement ->
  Rational {- ^ offset -} ->
  Maybe (Map Char (Maybe Location))

-- moving into a wall, not possible
moveBlock world _ _ loc _ _
  | isWall world loc = Nothing

-- move introduced a loop, trim off the tail and report success
moveBlock _ _ visited loc _ _
  | Just (n,_,_) <- Map.lookup loc visited
  = Just (Map.fromList [(c,l) | (a,c,l) <- Map.elems visited, a >= n])

moveBlock world locMap visited loc dir offset =
  case [(n,box) | (n,box) <- Map.assocs (worldBoxes world), boxLocation box == Just loc] of
    -- moving an empty space, so we're done
    [] -> Just (Map.fromList [(c,l) | (_,c,l) <- Map.elems visited])

    -- moving a box
    (name,box):_ ->
      case nextLoc world dir loc offset of
        -- block is exiting into the void; finish the move
        Nothing -> Just (Map.fromList ((name, Nothing) : [(c,l) | (_,c,l) <- Map.elems visited]))

        Just (loc', offset') ->
         do guard (not (isWall world loc'))
            moveBlock' world locMap visited loc loc' dir name box Set.empty offset'

moveBlock' ::
    World ->
    Map Location (Char, Box) ->
    Map Location (Int, Char, Maybe Location) ->
    Location ->
    Location ->
    Movement ->
    Char ->
    Box ->
    Set Location ->
    Rational {- ^ offset -} ->
    Maybe (Map Char (Maybe Location))
moveBlock' world locMap visited loc loc' dir name box enters offset =
        msum [moveTo, moveInto, moveToEat]
      where
        moveTo =
         do moveBlock world locMap (addVisited name loc (Just loc') visited) loc' dir 0

        moveInto =
         do (n,b) <- Map.lookup loc' locMap
            (locI, offset') <- enterLoc world n b dir offset
            if Set.member locI enters then
              moveEpsilon
            else
              moveBlock' world locMap visited loc locI dir name box (Set.insert locI enters) offset'

        moveEpsilon =
         do epsilon <- findEpsilon world (locName loc')
            let eBox = boxIx world epsilon
            (locI, offset') <- enterLoc world epsilon eBox dir offset
            moveBlock' world locMap visited loc locI dir name box (Set.insert locI enters) offset'

        moveToEat =
         do let dir' = invert dir
            (locE, _) <- enterLoc world name box dir' 0
            (name', box') <- Map.lookup loc' locMap
            moveBlock' world locMap (addVisited name loc (Just loc') visited) loc' locE dir' name' box' Set.empty 0

enterLoc :: World -> Char -> Box -> Movement -> Rational -> Maybe (Location, Rational)
enterLoc world name box dir@(dy,dx) offset =
 do name' <-
      case boxType box of
        Link c     -> Just c
        Original{} -> Just name
        Infinity{} -> Nothing
        Epsilon {} -> Just name

    let (bh, bw) = boxSize world box

        -- moving on y axis crosses width face and vice versa
        faceSize = abs dy * bw + abs dx * bh

        go y x = Just
          (Location name' y x,
           fromIntegral faceSize * offset
           - fromIntegral ((abs dy *x + abs dx*y)))

        ((ylo,xlo),(yhi,xhi)) = bounds (boxWalls world box)

    case dir of
      (-1, 0) -> go yhi (midpoint xlo xhi offset)
      ( 1, 0) -> go ylo (midpoint xlo xhi offset)
      ( 0,-1) -> go (midpoint ylo yhi offset) xhi
      ( 0, 1) -> go (midpoint ylo yhi offset) xlo
      _ -> error "enterLoc: bad direction"

boxIx :: HasCallStack => World -> Char -> Box
boxIx world name =
  case Map.lookup name (worldBoxes world) of
    Nothing -> error ("No such box: " ++ [name])
    Just box -> box

invert :: Movement -> Movement
invert (dy,dx) = (-dy, -dx)

midpoint :: Int -> Int -> Rational -> Int
midpoint lo hi offset = round (offset * fromIntegral (hi-lo+1))

addVisited ::
    Char {- ^ name -} ->
    Location {- ^ start -} ->
    Maybe Location {- ^ end   -} ->
    Map Location (Int, Char, Maybe Location) ->
    Map Location (Int, Char, Maybe Location)
addVisited name k v m = Map.insert k (Map.size m, name, v) m

nextLoc :: World -> Movement -> Location -> Rational -> Maybe (Location, Rational)
nextLoc world (dy, dx) = go Set.empty
  where

  go _ (Location b y x) offset
    | Just box <- Map.lookup b (worldBoxes world)
    , inRange (bounds (boxWalls world box)) (y+dy, x+dx)
    = Just (Location b (y+dy) (x+dx), offset)

  go visited (Location b y x) offset
    | Just box <- Map.lookup b (worldBoxes world)
    , Set.notMember b visited
    , Just boxLoc <- boxLocation box
    , let (bh,bw) = boxSize world box
    , let faceSize = abs dy * bw + abs dx * bh
    = go (Set.insert b visited) boxLoc
    $ (offset + fromIntegral (abs dy*x+abs dx*y))
    / fromIntegral faceSize

  go visited (Location b y x) offset
    | Set.member b visited
    , Just b' <- findInfinity world b
    = go visited (Location b' y x) offset

  go _ _ _ = Nothing

findInfinity :: World -> Char -> Maybe Char
findInfinity world b =
  listToMaybe [b' | (b', box) <- Map.assocs (worldBoxes world)
                  , Infinity i <- [boxType box], i == b]

findEpsilon :: World -> Char -> Maybe Char
findEpsilon world b =
  listToMaybe [b' | (b', box) <- Map.assocs (worldBoxes world)
                  , Epsilon i _ <- [boxType box], i == b]

worldLocations :: World -> Map Location (Char, Box)
worldLocations world =
  Map.fromList
    [(loc, (n,box))
       | (n, box) <- Map.assocs (worldBoxes world)
       , loc <- maybeToList (boxLocation box)]