{-# LANGUAGE OverloadedStrings #-}

module Encoding.Tap
  ( enctap
  , dectap
  ) where

import Data.Char (isAsciiUpper, isAsciiLower, toLower, chr, ord, toUpper)
import Data.List (find)
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as BC

enctap :: ByteString -> ByteString
enctap = BC.concatMap encodeChar
  where
    encodeChar ' ' = BC.singleton ' '
    encodeChar 'K' = BC.pack "22"  -- Special case: 'K' is encoded as "22"
    encodeChar c
      | isAsciiUpper c || isAsciiLower c =
          let (row, col) = positionInGrid c
          in BC.pack [chr (ord '0' + row), chr (ord '0' + col)]
      | otherwise = BC.empty  -- Handle non-alphabet characters or errors
    
    positionInGrid c =
      let idx = ord (toUpper c) - ord 'A'
          row = idx `div` 5 + 1
          col = idx `mod` 5 + 1
      in (row, col)

dectap :: ByteString -> ByteString
dectap = BC.concat . map decodePair . chunkPairs . BC.filter (/= ' ')
  where
    chunkPairs :: ByteString -> [ByteString]
    chunkPairs bs
      | BC.null bs = []
      | otherwise = let (pair, rest) = BC.splitAt 2 bs in pair : chunkPairs rest

    decodePair pair
      | pair == "22" = BC.singleton 'K'  -- Special case: "22" is decoded as 'K'
      | BC.length pair == 2 =
          let [row, col] = BC.unpack pair
              rowIdx = ord row - ord '0'
              colIdx = ord col - ord '0'
              idx = (rowIdx - 1) * 5 + (colIdx - 1)
          in BC.singleton $ chr (ord 'A' + idx)
      | otherwise = error $ "Invalid Tap Code sequence: " ++ show pair