import Data.Vect
import Data.List1
import Data.SnocList
import Data.String
import Data.Fin
import System.File.ReadWrite

allFins' : (n : Nat) -> List (Fin n)
allFins' n = toList (allFins'' n)
  where
    allFins'' : (n : Nat) -> SnocList (Fin n)
    allFins'' 0 = Lin
    allFins'' 1 = Lin :< FZ
    allFins'' (S k) = map weaken (allFins'' k) :< last

data Data : Type where
  -- An atom, in this case, a nat, for now
  MkAtom : (atom : Nat) -> Data
  -- A list of sub-datas
  MkList : (n : Nat) -> (contents : Vect n Data) -> Data

Show Data where
  show (MkAtom atom) = show atom
  show input@(MkList n contents) =
    let inner = joinBy "," . toList . map (\x => show (assert_smaller input x)) $ contents in
    "[" ++ inner ++ "]"

atomToList : Nat -> Data
atomToList i = MkList 1 [MkAtom i]

splitNonNesting : String -> Maybe (List String)
splitNonNesting str =
  do result <- splitNonNesting' (== ',') (unpack (trim str)) 0 Lin
     let packed = map pack result
     pure packed
  where splitNonNesting' : (test : (Char -> Bool)) -> (input: List Char) -> (depth : Nat)
          -> (buffer : SnocList Char) -> Maybe (List (List Char))
        splitNonNesting' test (['[', ']']) depth buffer = Just [[]]
        splitNonNesting' test ('[' :: xs) 0 buffer =
          splitNonNesting' test xs 1 buffer
        splitNonNesting' test ('[' :: xs) depth@(S k) buffer =
          splitNonNesting' test xs (S depth) (buffer :< '[')
        splitNonNesting' test ([']']) 0 [<] = Just []
        splitNonNesting' test ([']']) 0 buffer = Nothing
        splitNonNesting' test ([']']) 1 buffer@(sx :< x) =
          let buffer = toList buffer in Just [buffer]
        splitNonNesting' test ([']']) (S k) buffer = Nothing
        splitNonNesting' test [] 0 [<] = Just []
        splitNonNesting' test [] 0 buffer@(sx :< x) =
          let buffer = toList buffer in Just [buffer]
        splitNonNesting' test [] (S k) buffer = Nothing
        splitNonNesting' test (']' :: xs) 0 buffer = Nothing
        splitNonNesting' test (']' :: xs) (S k) buffer =
          splitNonNesting' test xs k (buffer :< ']')
        splitNonNesting' test (x :: xs) 0 buffer = Nothing
        splitNonNesting' test (x :: xs) 1 buffer =
          if test x
          then do let buffer = toList buffer
                  rest <- splitNonNesting' test xs 1 Lin
                  pure $ buffer :: rest
          else splitNonNesting' test xs 1 (buffer :< x)
        splitNonNesting' test (x :: xs) (S k) buffer =
          splitNonNesting' test xs (S k) (buffer :< x)

parseAtom : String -> Maybe Data
parseAtom = map MkAtom . parsePositive

mutual
  parseList : String -> Maybe Data
  parseList str =
    do components <- splitNonNesting str
       parsed <- traverse parseData components
       let n = length parsed
       list <- toVect n parsed
       pure (MkList n list)

  parseData : String -> Maybe Data
  parseData str =
    -- Attempt to parse as an atom first, then a list
    case parseAtom str of
      Just atom => Just atom
      Nothing =>
        -- Then attempt to parse as a list
        parseList str

checkOrder' : (left : Data) -> (right : Data) -> Maybe Bool
checkOrder' (MkAtom atom) (MkAtom k) = if atom == k then Nothing else Just (atom < k)
checkOrder' (MkAtom atom) list@(MkList n contents) = checkOrder' (atomToList atom) list
checkOrder' list@(MkList n contents) (MkAtom atom) = checkOrder' list (atomToList atom)
checkOrder' (MkList 0 []) (MkList 0 []) = Nothing
checkOrder' (MkList 0 []) (MkList m ys) = Just True
checkOrder' (MkList (S n) xs) (MkList 0 []) = Just False
checkOrder' (MkList (S n) (x :: xs)) (MkList (S m) (y :: ys)) =
  case checkOrder' x y of
    Nothing => checkOrder' (MkList n xs) (MkList m ys)
    Just z => Just z
-- For some reason idris thinks this isn't covering
checkOrder' _ _ = Nothing

-- Returns true of the two datas are in the correct order
checkOrder : (left : Data) -> (right : Data) -> Bool
checkOrder left right =
  case checkOrder' left right of
    Nothing => True
    Just x => x

parsePairs : String -> Maybe (n : Nat ** Vect n (Data, Data))
parsePairs str =
  do let lines = map trim . lines $ str
     let groups = split (=="") lines
     pairs <- traverse toPair (toList groups)
     pairs <- traverse handlePairs pairs
     let n = length pairs
     pairs <- toVect n pairs
     pure (n ** pairs)
  where toPair : List String -> Maybe (String, String)
        toPair [a, b] = Just (a, b)
        toPair strs = Nothing
        handlePairs : (String, String) -> Maybe (Data, Data)
        handlePairs (x, y) =
          do x <- parseData x
             y <- parseData y
             pure (x, y)

simple = """
  [1,1,3,1,1]
  [1,1,5,1,1]

  [[1],[2,3,4]]
  [[1],4]

  [9]
  [[8,7,6]]

  [[4,4],4,4]
  [[4,4],4,4,4]

  [7,7,7,7]
  [7,7,7]

  []
  [3]

  [[[]]]
  [[]]

  [1,[2,[3,[4,[5,6,7]]]],8,9]
  [1,[2,[3,[4,[5,6,0]]]],8,9]
  """

testPart1 : IO ()
testPart1 =
  do Just (n ** pairs) <- pure (parsePairs simple)
       | Nothing => putStrLn "Failed to parse pairs"
     let indicies = allFins' n
     let correct_order =
       map S .
       map finToNat .
       filter (\x =>
                let (x,y) = index x pairs
                in checkOrder x y) $
       indicies
     printLn pairs
     printLn correct_order

part1 : (pairs : (n : Nat ** Vect n (Data, Data))) -> Nat
part1 (n ** pairs) =
  let indicies = allFins' n
      correct_order =
        map S .
        map finToNat .
        filter (\x =>
                 let (x,y) = index x pairs
                 in checkOrder x y) $
        indicies
  in sum correct_order

main : IO ()
main =
  do Right file <- readFile "input"
       | Left err => printLn err
     Just pairs <- pure (parsePairs file)
       | Nothing => putStrLn "Failed parsing input"
     putStrLn $ "Part 1: " ++ show (part1 pairs)