import Data.Vect
import Data.List
import Data.Fin
import System.File.ReadWrite

-- Get windows of a specific size over the data
windows : {n : Nat} -> List a -> List (Vect n a)
windows [] = []
windows {n} xs =
  let (head, _) = splitAt n xs
      tail = drop 1 xs
  in case toVect n head of
      Nothing => []
      (Just x) => x :: windows tail

findStart : Nat -> List Char -> Maybe Nat
findStart n cs =
  let segments = windows {n} cs
  in case findIndex (\x => let ( p ** _ ) = nub x in p == n) segments of
       Nothing => Nothing
       (Just x) => Just $ finToNat x + n

findStartOfPacket : List Char -> Maybe Nat
findStartOfPacket = findStart 4

findStartOfMessage : List Char -> Maybe Nat
findStartOfMessage = findStart 14

testCases : List (String, Nat, Nat)
testCases = [("mjqjpqmgbljsphdztnvjfqwrcgsmlb", 7, 19),
             ("bvwbjplbgvbhsrlpgdmjqwftvncz", 5, 23),
             ("nppdvjthqldpwncqszvftbrmjlhg", 6, 23),
             ("nznrnfrfntjfmvfwmzdfjlvtqnbhcprsg", 10, 29),
             ("zcfzfwzzqfrljwzlrfnpqdbhtmscgvjw", 11, 26)]

testCase : (String, Nat, Nat) -> IO Bool
testCase (string, start, message) =
  do putStrLn $ "Test Case: " ++ string
     putStrLn $ "Expected Start of packet: " ++ show start
     putStrLn $ "Expected Start of message: " ++ show message
     let result = findStartOfPacket (unpack string)
     let messageResult = findStartOfMessage (unpack string)
     putStrLn $ "Evaluated start of packet: " ++ show result
     putStrLn $ "Evaluated start of message: " ++ show messageResult
     let result = not (result == Just start)
     result <- if result
               then do putStrLn "Test Failed!!!"
                       pure True
               else if messageResult == Just message
                      then do putStrLn "Test Passed"
                              pure False
                      else do putStrLn "Test Failed!!!"
                              pure True
     putStrLn ""
     pure result

main : IO ()
main =
  do testsPart1 <- traverse testCase testCases
     if or . map (\x => Delay x) $ testsPart1
       then putStrLn "Failed tests"
       else
         do file <- readFile "input"
            case file of
              Right contents =>
                do putStr "Part 1: "
                   putStrLn $ case findStartOfPacket (unpack contents) of
                               Just start => show start
                               Nothing => "Failed to find start"
                   putStr "Part 2: "
                   putStrLn $ case findStartOfMessage (unpack contents) of
                               Just start => show start
                               Nothing => "Failed to find start"

              Left err => printLn err