{-# LANGUAGE OverloadedStrings #-}

module Main where

import qualified Data.ByteString.Lazy as BL
import Data.Char (isSpace)
import qualified Data.Csv as CSV
import Data.List (delete, dropWhileEnd, minimumBy, transpose)
import Data.Maybe (catMaybes, isNothing, listToMaybe)
import Data.Ord (comparing)
import qualified Data.Vector as V
import System.Environment (getArgs)
import Text.Read (readMaybe)

type Cell = Maybe Int
type Grid = [[Cell]]
type Pos = (Int, Int)

parseCell :: String -> Cell
parseCell s = readMaybe (trim s)
  where
    trim = dropWhileEnd isSpace . dropWhile isSpace

readCSV :: FilePath -> IO Grid
readCSV path = do
    csvData <- BL.readFile path
    case CSV.decode CSV.NoHeader csvData of
        Left err -> error ("CSV parse error: " ++ err)
        Right v -> return $ map (map (parseCell . trimQuotes) . V.toList) (V.toList v)
  where
    trimQuotes = filter (/= '"')

printGrid :: Grid -> IO ()
printGrid = putStrLn . unlines . map (unwords . map showCell)
  where
    showCell = maybe "." show

solve :: Grid -> Maybe Grid
solve g
    | isComplete g = Just g
    | otherwise = case nextCell g of
        Nothing -> Nothing
        Just (r, c, opts) ->
            listToMaybe
                [ result
                | n <- opts
                , let g' = updateGrid g r c (Just n)
                , Just result <- [solve g']
                , isComplete result
                ]

isComplete :: Grid -> Bool
isComplete = all (notElem Nothing)

nextCell :: Grid -> Maybe (Int, Int, [Int])
nextCell g =
    case filter (not . null . (\(_, _, opts) -> opts)) candidates of
        [] -> Nothing
        cs -> Just $ minimumBy (comparing (\(_, _, opts) -> length opts)) cs
  where
    candidates =
        [ (r, c, validOptions g r c)
        | (r, row) <- zip [0 ..] g
        , (c, cell) <- zip [0 ..] row
        , isNothing cell
        ]

validOptions :: Grid -> Int -> Int -> [Int]
validOptions g r c =
    foldr delete [1 .. 9] (catMaybes (getRow r g ++ getCol c g ++ getBox r c g))

getRow :: Int -> Grid -> [Cell]
getRow r g = g !! r

getCol :: Int -> Grid -> [Cell]
getCol c = map (!! c)

getBox :: Int -> Int -> Grid -> [Cell]
getBox r c g =
    [ g !! r' !! c'
    | r' <- [br .. br + 2]
    , c' <- [bc .. bc + 2]
    ]
  where
    br = r - r `mod` 3
    bc = c - c `mod` 3

updateGrid :: Grid -> Int -> Int -> Cell -> Grid
updateGrid g r c val =
    take r g
        ++ [take c (g !! r) ++ [val] ++ drop (c + 1) (g !! r)]
        ++ drop (r + 1) g

main :: IO ()
main = do
    args <- getArgs
    case args of
        [filePath] -> do
            grid <- readCSV filePath
            case solve grid of
                Nothing -> putStrLn "No complete solution found."
                Just g -> printGrid g
        _ -> putStrLn "Usage: sudoku-solver <sudoku.csv>"
