module Calendar where

import Control.Monad (forM)
import Control.Monad.Trans.State.Strict (evalState, state)
import Data.Maybe (isNothing)
import Data.Time

import Util


-- | The list of @a@ values must be as long as the number of days between the
-- two bounds, which are inclusive. Values are returned with the appropriate
-- day in the table.
--
-- * Returns a list of years with:
--   * a list of rows, with in each row
--     * a list of months (empty spacer if Nothing), with inside each month
--       * a list of weeks (rows), with in each week
--         * a list of dates in the week (empty spacer if Nothing).
--
-- >>> calendarLayout (fromGregorian 2025 12 16) (fromGregorian 2026 2 3) (cycle ['A'..'Z'])
calendarLayout :: Day -> Day -> [a] -> [(Year, [[Maybe (MonthOfYear, [[Maybe (DayOfMonth, a)]])]])]
calendarLayout startDay endDay values =
  let (startYear, startMonth, startDate) = toGregorian startDay
      (endYear, endMonth, endDate) = toGregorian endDay
      takeValue = state $ \case val:vals -> (val, vals)
                                [] -> error "calendarLayout: too few values"
  in flip evalState values $
       forM [startYear .. endYear] $ \year -> do
         let row1 | year == startYear = (startMonth - 1) `div` 4
                  | otherwise = 0
             row2 | year == endYear = (endMonth - 1) `div` 4
                  | otherwise = 2
         fmap (year,) . forM [row1..row2] $ \monRowIdx -> do
           let mon1 | year == startYear && monRowIdx == row1 = startMonth
                    | otherwise = 4 * monRowIdx + 1
               mon2 | year == endYear && monRowIdx == row2 = endMonth
                    | otherwise = 4 * monRowIdx + 4
           forM [4 * monRowIdx + 1 .. 4 * monRowIdx + 4] $ \month ->
             if month < mon1 || month > mon2 then return Nothing else fmap Just $ do
               let cal = monthCalendar year month
                   filterCal predicate = map (map (\md -> if maybe False predicate md then md else Nothing))
                   cal' | year == startYear && monRowIdx == row1 && month == mon1 = filterCal (>= startDate) cal
                        | year == endYear   && monRowIdx == row2 && month == mon2 = filterCal (<= endDate)   cal
                        | otherwise = cal
                   cal'' = dropWhileEnd null $ map (\week -> if all isNothing week then [] else week) cal'
               fmap (month,) . forM cal'' $ \week ->
                 forM week $ \mdate ->
                   traverse (\d -> (d,) <$> takeValue) mdate

-- >>> import Text.Printf
-- >>> error $ unlines . map (unwords . map (maybe "  " (printf "%2d"))) $ monthCalendar 2026 4
--        1  2  3  4  5
--  6  7  8  9 10 11 12
-- 13 14 15 16 17 18 19
-- 20 21 22 23 24 25 26
-- 27 28 29 30
monthCalendar :: Year -> MonthOfYear -> [[Maybe DayOfMonth]]
monthCalendar year month
  | [1, 7] <- map fromEnum [Monday, Sunday] =
      let wkday1 = dayOfWeek (fromGregorian year month 1)
          nskipsRow1 = fromEnum wkday1 - 1
          lastDate = gregorianMonthLength year month
          genRows date =
            map Just [date .. min (date + 6) lastDate] : (if lastDate <= date + 6 then [] else genRows (date + 7))
      in (replicate nskipsRow1 Nothing ++ map Just [1 .. 7 - nskipsRow1])
         : genRows (7 - nskipsRow1 + 1)

  | otherwise = error "DayOfWeek enum is unexpected"