module Main where

import Data.Char (digitToInt)
import Data.Maybe (fromJust, fromMaybe, isJust, mapMaybe)
import qualified Data.Vector as V
import Types
  ( Grid (Grid),
    allCells,
    colView,
    count,
    getCell,
    initGrid,
    mapGrid,
    printGrid,
    putCell,
    rowView,
    squareView,
  )

main :: IO ()
main = do
  putStr $ printSudoku dummyGrid
  print $ isValidGrid dummyGrid
  putStr . printGrid . fromJust $ solve dummyGrid

test :: String
test = "070000043040009610800634900094052000358460020000800530080070091902100005007040802"

dummyGrid :: Grid (Maybe Int)
dummyGrid = initGrid 3 (map toMaybeInt test)
  where
    toMaybeInt :: Char -> Maybe Int
    toMaybeInt c = if int == 0 then Nothing else Just int
      where
        int = digitToInt c

printSudoku :: Grid (Maybe Int) -> String
printSudoku gr = printGrid mappedGrid
  where
    mappedGrid :: Grid Int
    mappedGrid = mapGrid (fromMaybe 0) gr

isValidGrid :: Grid (Maybe Int) -> Bool
isValidGrid gr@(Grid _ d) =
  all isValidRow [0 .. d * d - 1]
    && all isValidCol [0 .. d * d - 1]
    && all isValidSquare [0 .. d * d - 1]
  where
    validView :: V.Vector (Maybe Int) -> Bool
    validView v = V.foldl (\drag (e, c) -> case e of Nothing -> drag; Just _ -> (c < 2) && drag) True counts
      where
        counts = V.map (\e -> (e, count v e)) v
    isValidRow :: Int -> Bool
    isValidRow y = validView $ rowView gr y
    isValidCol :: Int -> Bool
    isValidCol x = validView $ colView gr x
    isValidSquare :: Int -> Bool
    isValidSquare n = validView $ squareView gr n

isFull :: Grid (Maybe Int) -> Bool
isFull gr = allCells gr isJust

isSolved :: Grid (Maybe Int) -> Bool
isSolved gr = isValidGrid gr && isFull gr

extractSolution :: Grid (Maybe Int) -> Grid Int
extractSolution gr =
  if not (isSolved gr)
    then error "trying to extract solution from unsolved grid"
    else mapGrid (fromMaybe 0) gr

solveHelper :: Grid (Maybe Int) -> (Int, Int) -> Maybe (Grid Int)
solveHelper gr@(Grid _ d) (x, y)
  | not (isValidGrid gr) = Nothing
  | isSolved gr = Just (extractSolution gr)
  | otherwise = case cell of
      Nothing -> if not (null solvedGrids) then head solvedGrids else Nothing
      Just _ -> solveHelper gr nextCoord
  where
    placeAndSolve :: Int -> Maybe (Grid Int)
    placeAndSolve n = solveHelper newGrid nextCoord
      where
        newGrid = putCell gr (x, y) (Just n)
    solvedGrids = map Just (mapMaybe placeAndSolve [1 .. d * d])
    cell = getCell gr (x, y)
    nextCoord = if x < d * d - 1 then (x + 1, y) else (0, y + 1)

solve :: Grid (Maybe Int) -> Maybe (Grid Int)
solve gr = solveHelper gr (0, 0)