{-# Language ImportQualifiedPost, BlockArguments, LambdaCase #-}
module Main where

import Data.Array
import Data.Map (Map)
import Data.Map qualified as Map
import Graphics.Vty
import Control.Exception
import Data.List (intersperse)
import Data.List.NonEmpty (NonEmpty(..))
import Data.List.NonEmpty qualified as NonEmpty
import Debug.Trace
import Control.Monad
import Data.Maybe
import Data.Set (Set)
import Data.Set qualified as Set

type Coord = (Int, Int)

data Box = Box {
  boxLocation :: Location,
  boxWalls    :: Array Coord Bool,
  boxColor    :: Attr
}
   deriving (Show, Read, Eq)

data Location = Location Char Int Int
  deriving (Read, Show, Ord, Eq)

data World = World {
  worldBoxes :: Map Char Box,
  worldMe :: Char
}
  deriving (Show, Read, Eq)

renderBox :: World -> Map Location Char -> Box -> Char -> Image
renderBox world locMap box boxName =
  vertCat
  [
    horizCat
    [
      if boxWalls box ! (y,x) then unit myAttr '▓'
      else case Map.lookup (Location boxName y x) locMap of
        Nothing -> unit myAttr '░'
        Just n -> unit (boxColor (worldBoxes world Map.! n)) n
      | x <- [xlo .. xhi]
    ]

    | let ((ylo,xlo),(yhi,xhi)) = bounds (boxWalls box)
    , let myAttr = boxColor box
    , y <- [ylo .. yhi]
    , let unit a x = string a [x,x]
  ]

renderBox' :: World -> Map Location Char -> Box -> Char -> ((Int,Int),(Int,Int)) -> Int -> Image
renderBox' world locMap box boxName ((ylo,xlo),(yhi,xhi)) scale
  | scale == 1 = renderBox world locMap box boxName
  | otherwise =
  vertCat
  [
    horizCat
    [
      drawAt boxName box y x | x <- [xlo .. xhi]
    ]
    | y <- [ylo .. yhi]
  ]
  where
    unit a x =
      vertCat (replicate scale (string a (replicate (2*scale) x)))
    drawAt boxName box y x
      | let goodCoord = inRange (bounds (boxWalls box))
      , not (goodCoord (y,x))
      = unit defAttr '?'
      | otherwise
      , let myAttr = boxColor box =
        if boxWalls box ! (y,x) then unit myAttr '▓'
        else case Map.lookup (Location boxName y x) locMap of
          Nothing -> unit myAttr '░'
          Just n -> renderBox' world locMap box' n (bounds (boxWalls box')) (scale `div` boxSize box)
            where
              box' = worldBoxes world Map.! n

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 n = (- (n-1)`div`2, n`div`2)

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

drawNestedWorld :: World -> Image
drawNestedWorld world = renderBox' world locMap box0 name0 bnds 49
  where
    locMap = Map.fromList [(boxLocation loc, n) | (n, loc) <- Map.toList (worldBoxes world)]
    (bnds,name0) = go 2 0 0 (worldMe world)
    box0 = worldBoxes world Map.! name0
    
    go 0 y x name = (((y-1,x-1),(y+1,x+1)), name)
    go n y x name =
      case Map.lookup name (worldBoxes world) of
        Nothing -> (((y-1,x-1),(y+1,x+1)), name)
        Just box -> go (n-1) y x name'
          where
            Location name' y x = boxLocation box




drawWorld :: World -> Image
drawWorld world =
  horizCat $
  intersperse (char defAttr ' ')
  [
    renderBox world locMap box boxName
    | (boxName, box) <- Map.toList (worldBoxes world)
  ]
  where
    locMap = Map.fromList [(boxLocation loc, n) | (n, loc) <- Map.toList (worldBoxes world)]

solid n = makeWalls (replicate n (replicate n 'x'))

world0 :: World
world0 = World {
  worldMe = 'b',
  worldBoxes = Map.fromList
    [('1', Box {
      boxColor = withForeColor defAttr green,
      boxLocation = Location '1' (-1) 0,
      boxWalls = makeWalls [
        "▓▓   ▓▓",
        "▓     ▓",
        "▓     ▓",
        "▓     ▓",
        "▓  ▓  ▓",
        "▓     ▓",
        "▓  ▓  ▓"
        ]
    }),
    ('2', Box {
      boxColor = withForeColor defAttr cyan,
      boxLocation = Location '1' 1 1,
      boxWalls = makeWalls [
        "▓▓▓ ▓▓▓",
        "▓     ▓",
        "▓     ▓",
        "       ",
        "▓     ▓",
        "▓     ▓",
        "▓▓▓ ▓▓▓"
      ]
    }),
    ('3', Box {
      boxColor = withForeColor defAttr blue,
      boxLocation = Location '2' 1 1,
      boxWalls = makeWalls [
        "▓▓   ▓▓",
        "▓  ▓  ▓",
        "▓     ▓",
        "▓     ▓",
        "▓     ▓",
        "▓     ▓",
        "▓▓▓ ▓▓▓"
      ]
    }),
        ('4', Box {
      boxColor = withForeColor defAttr black,
      boxLocation = Location '3' (-3) 0,
      boxWalls = makeWalls [
        "▓▓   ▓▓",
        "▓  ▓  ▓",
        "▓     ▓",
        "▓     ▓",
        "▓     ▓",
        "▓     ▓",
        "▓▓▓ ▓▓▓"
      ]
    }),
    ('b', Box {
        boxColor = withForeColor defAttr red,
        boxLocation = Location '1' 0 1,
        boxWalls = solid 7
    }),
    ('x', Box {
        boxColor = withForeColor defAttr yellow,
        boxLocation = Location '1' 0 (-1),
        boxWalls = solid 7
    }),
    ('y', Box {
        boxColor = withForeColor defAttr magenta,
        boxLocation = Location '1' 0 (-2),
        boxWalls = solid 7
    })

    ]
  }

-- Move an object
-- 1. remove it from the world
-- 2. compute where it would move to
-- 3. a. that spot is empty
--    b. try to move that object forward
--    c. try to move that object backward into me

move :: World -> (Int,Int) -> World
move world dir =
  case moveBlock world Map.empty (myLocation world) dir 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 :: Location -> World -> Bool
isWall (Location n y x) world =
  case Map.lookup n (worldBoxes world) of
    Nothing -> True
    Just box -> boxWalls box ! (y,x)

main :: IO ()
main =
  bracket (mkVty =<< userConfig) shutdown \vty ->
   do
      loop vty (pure world0)
      pure ()

loop :: Vty -> NonEmpty World -> IO ()
loop vty (world :| history) =
 do update vty (picForImage (drawNestedWorld world))
    ev <- nextEvent vty
    case ev of
      EvKey key _modifier ->
        case key of
          KUp    -> loop vty (move world (-1,0) :| world : history)
          KDown  -> loop vty (move world (1,0)  :| world : history)
          KLeft  -> loop vty (move world (0,-1) :| world : history)
          KRight -> loop vty (move world (0,1)  :| world : history)
          KChar 'z'
            | Just worlds <- NonEmpty.nonEmpty history ->
              loop vty worlds
          KEsc -> pure ()
          _ -> loop vty (world :| history)
      _ -> loop vty (world :| history)

type Movement = (Int, Int)

moveBlock ::
  World ->
  Map Location (Int, Location) ->
  Location ->
  Movement ->
  Maybe (Map Location Location)

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

-- move introduced a loop, trim off the tail and report success
moveBlock world visited loc _
  | Just (n,_) <- Map.lookup loc visited
  = Just (fmap snd (Map.filter (\(a,_)->a >= n) visited))

moveBlock world visited loc dir =
  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' <- nextLoc world loc dir
        guard (not (isWall loc' world))
        moveBlock' world visited loc loc' dir name box

moveBlock' world visited loc loc' dir name box =
        msum [moveTo loc', moveInto loc', moveToEat loc']
      where
        moveTo loc' =
         do moveBlock world (addVisited loc loc' visited) loc' dir

        moveInto loc' =
         do (n,b) <- boxAt world loc'
            let locI = enterLoc n b dir
            moveBlock' world visited loc locI dir name box -- beware epsilon!
        
        moveToEat loc' =
         do let dir' = invert dir
            let locE = enterLoc name box dir'
            (name', box') <- boxAt world loc' 
            moveBlock' world (addVisited loc loc' visited) loc' locE dir' name' box' 

enterLoc :: Char -> Box -> Movement -> Location
enterLoc name box dir =
  case dir of
    (-1, 0) -> Location name yhi (midpoint xlo xhi)
    ( 1, 0) -> Location name ylo (midpoint xlo xhi)
    ( 0,-1) -> Location name (midpoint ylo yhi) xhi
    ( 0, 1) -> Location name (midpoint ylo yhi) xlo
    _ -> error "enterLoc: bad direction"
  where
    ((ylo,xlo),(yhi,xhi)) = bounds (boxWalls box)

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

invert (dy,dx) = (-dy, -dx)

midpoint lo hi = (hi+lo)`div`2

addVisited k v m = Map.insert k (Map.size m, v) m

nextLoc world loc (dy, dx) = go Set.empty loc
  where
    
  go visited (Location b y x)
    | Just box <- Map.lookup b (worldBoxes world)
    , inRange (bounds (boxWalls box)) (y+dy, x+dx)
    = Just (Location b (y+dy) (x+dx))
  
  go visited (Location b y x)
    | Just box <- Map.lookup b (worldBoxes world)
    , Set.notMember b visited
    = go (Set.insert b visited) (boxLocation box)

  go _ _ = Nothing