module Model where

import Data.Array
import Data.Map (Map)
import Data.Map qualified as Map
import Graphics.Vty
import Control.Monad
import Data.Maybe
import Data.Set qualified as Set
import Data.Set (Set)

type Coord = (Int, Int)

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

boxWalls :: World -> Box -> Array Coord Bool
boxWalls world box =
    case boxType box of
        Original walls -> walls
        Link c -> boxWalls world (worldBoxes world Map.! c)
        Infinity c -> boxWalls world (worldBoxes world Map.! c)
        Epsilon _ walls -> walls

data BoxType
    = Original (Array Coord Bool)
    | Link     Char
    | Infinity Char
    | Epsilon  Char (Array Coord Bool)
   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
}
  deriving (Show, Read, Eq)

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

makeWalls :: [String] -> Array Coord Bool
makeWalls rows = listArray ((ylo,xlo),(yhi,xhi)) (map (' '/=) (concat rows))
  where
    h = length rows
    w = length (head rows)
    (xlo,xhi) = mkRange w
    (ylo,yhi) = mkRange h

mkRange :: Int -> (Int,Int)
mkRange n = (- (n-1)`div`2, n`div`2)

boxSize :: World -> Box -> Int
boxSize world box = yhi-ylo+1
  where
    ((ylo,_),(yhi,_)) = bounds (boxWalls world box)


solid :: Int -> Array Coord Bool
solid n = makeWalls (replicate n (replicate n 'x'))

move :: World -> (Int,Int) -> World
move world dir =
  case moveBlock world Map.empty (myLocation world) dir 0 of
    Nothing -> world
    Just changes ->
      let f box = box { boxLocation = Map.findWithDefault (boxLocation box) (boxLocation box) changes } in
      world { worldBoxes = fmap f (worldBoxes world)}

myLocation :: World -> Location
myLocation world =
  boxLocation (worldBoxes world Map.! 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 (Int, Location) ->
  Location ->
  Movement ->
  Rational {- ^ offset -} ->
  Maybe (Map Location 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 (fmap snd (Map.filter (\(a,_)->a >= n) visited))

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

    -- moving a box
    (name,box):_ ->
     do (loc', offset') <- nextLoc world dir loc offset
        --traceM (show "Next loc " ++ show loc' ++ " offset " ++ show offset')
        guard (not (isWall world loc'))
        moveBlock' world visited loc loc' dir name box Set.empty offset'

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

        moveInto =
         do (n,b) <- boxAt world loc'
            (locI, offset') <- enterLoc world n b dir offset
            if Set.member locI enters then do
              epsilon <- findEpsilon world (locName loc')
              let eBox = worldBoxes world Map.! epsilon
              (locI, offset') <- enterLoc world epsilon eBox dir offset
              moveBlock' world visited loc locI dir name box (Set.insert locI enters) offset'
            else
              moveBlock' world 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') <- boxAt world loc' 
            moveBlock' world (addVisited loc 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 go y x = Just
          (Location name' y x,
           fromIntegral(boxSize world box) * 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"

boxAt :: World -> Location -> Maybe (Char, Box)
boxAt world loc =
  listToMaybe  [(n,b) | (n,b) <- Map.assocs (worldBoxes world), boxLocation b == loc]

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

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

addVisited ::
    Location {- ^ start -} ->
    Location {- ^ end   -} ->
    Map Location (Int, Location) ->
    Map Location (Int, Location)
addVisited k v m = Map.insert k (Map.size m, 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
    = go (Set.insert b visited) (boxLocation box)
    $ (offset + fromIntegral (abs dy*x+abs dx*y))
      / fromIntegral (boxSize world box)

  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]