module ContentStore where

import Research
import Parser

import Control.Monad                  (foldM, forM)
import Data.ByteString                (ByteString)
import Data.List                      (nub, sort)
import Data.Maybe                     (catMaybes, fromJust)
import Data.Text                      (Text)
import Database.SQLite.Simple
import Database.SQLite.Simple.FromRow (FromRow(..), field)
import System.Directory               (createDirectoryIfMissing, getXdgDirectory, XdgDirectory(..))
import System.FilePath                ((</>), takeDirectory)

import qualified Data.ByteString      as BS
import qualified Data.Map             as Map
import qualified Data.Text            as T

data StoredNode = StoredNode ByteString deriving (Show)

instance FromRow StoredNode where
  fromRow = StoredNode <$> field

data StoredTerm = StoredTerm
  { termHash :: Text
  , termNames :: Text
  , termMetadata :: Text
  , termCreatedAt :: Integer
  , termTags :: Text
  } deriving (Show)

instance FromRow StoredTerm where
  fromRow = StoredTerm <$> field <*> field <*> field <*> field <*> field

parseNameList :: Text -> [Text]
parseNameList = filter (not . T.null) . T.splitOn ","

serializeNameList :: [Text] -> Text
serializeNameList = T.intercalate "," . nub . sort

initContentStore :: IO Connection
initContentStore = do
  dbPath <- getContentStorePath
  createDirectoryIfMissing True (takeDirectory dbPath)
  conn <- open dbPath
  execute_ conn "CREATE TABLE IF NOT EXISTS terms (\
                \hash TEXT PRIMARY KEY, \
                \names TEXT, \
                \term_data BLOB, \
                \metadata TEXT, \
                \created_at INTEGER DEFAULT (strftime('%s','now')), \
                \tags TEXT DEFAULT '')"
  execute_ conn "CREATE INDEX IF NOT EXISTS terms_names_idx ON terms(names)"
  execute_ conn "CREATE INDEX IF NOT EXISTS terms_tags_idx ON terms(tags)"
  execute_ conn "CREATE TABLE IF NOT EXISTS merkle_nodes (\
                \hash TEXT PRIMARY KEY, \
                \node_data BLOB NOT NULL)"
  return   conn

getContentStorePath :: IO FilePath
getContentStorePath = do
  dataDir <- getXdgDirectory XdgData "tricu"
  return $ dataDir </> "content-store.db"



hashTerm :: T -> Text
hashTerm = nodeHash . buildMerkle

storeTerm :: Connection -> [String] -> T -> IO Text
storeTerm conn newNamesStrList term = do
  let termHashText = hashTerm term
      newNamesTextList = map T.pack newNamesStrList
      metadataText = T.pack "{}"
  -- Store all Merkle nodes for this term
  _ <- storeMerkleNodes conn term
  existingNamesQuery <- query conn
    "SELECT names FROM terms WHERE hash = ?"
    (Only termHashText) :: IO [Only Text]

  case existingNamesQuery of
    [] -> do
      let allNamesToStore = serializeNameList newNamesTextList
      execute conn
        "INSERT INTO terms (hash, names, term_data, metadata, tags) VALUES (?, ?, ?, ?, ?)"
        (termHashText, allNamesToStore, BS.pack [], metadataText, T.pack "")
    [(Only currentNamesText)] -> do
      let currentNamesList = parseNameList currentNamesText
      let combinedNamesList = currentNamesList ++ newNamesTextList
      let allNamesToStore = serializeNameList combinedNamesList
      execute conn
        "UPDATE terms SET names = ?, metadata = ? WHERE hash = ?"
        (allNamesToStore, metadataText, termHashText)
    _ -> error $ "Multiple terms with same hash? " ++ show (length existingNamesQuery)

  return termHashText

-- | Reconstruct a Tree Calculus term from its Merkle root hash.
--   Recursively loads nodes and rebuilds the T structure.
loadTree conn h
  | h == nodeHash NLeaf = return (Just Leaf)  -- NLeaf is implicit, not stored
  | otherwise = do
      maybeNode <- getNodeMerkle conn h
      case maybeNode of
        Nothing -> return Nothing
        Just node -> Just <$> buildTree node
  where
    buildTree :: Node -> IO T
    buildTree (NStem childHash) = do
      child <- fromJust <$> loadTree conn childHash
      return (Stem child)
    buildTree (NFork lHash rHash) = do
      left <- fromJust <$> loadTree conn lHash
      right <- fromJust <$> loadTree conn rHash
      return (Fork left right)

-- | Store all nodes of a Merkle DAG by traversing the Term and building/storing nodes.
--   Returns the hash of the root node.
storeMerkleNodes :: Connection -> T -> IO MerkleHash
storeMerkleNodes _ Leaf = return $ nodeHash NLeaf
storeMerkleNodes conn (Stem t) = do
  childHash <- storeMerkleNodes conn t
  let thisNode = NStem childHash
  putMerkleNode conn thisNode
  return $ nodeHash thisNode
storeMerkleNodes conn (Fork l r) = do
  leftHash <- storeMerkleNodes conn l
  rightHash <- storeMerkleNodes conn r
  let thisNode = NFork leftHash rightHash
  putMerkleNode conn thisNode
  return $ nodeHash thisNode


-- | Insert a Merkle node into the store (idempotent).
putMerkleNode :: Connection -> Node -> IO ()
putMerkleNode conn node =
  execute conn "INSERT OR IGNORE INTO merkle_nodes (hash, node_data) VALUES (?, ?)"
    (nodeHash node, serializeNode node)

-- | Retrieve a Merkle node by its hash.
getNodeMerkle :: Connection -> MerkleHash -> IO (Maybe Node)
getNodeMerkle conn h =
  queryMaybeOne conn "SELECT node_data FROM merkle_nodes WHERE hash = ?" (Only h) >>= \case
    Just (StoredNode bs) -> return $ Just (deserializeNode bs)
    Nothing              -> return Nothing



hashToTerm :: Connection -> Text -> IO (Maybe StoredTerm)
hashToTerm conn hashText =
  queryMaybeOne conn (selectStoredTermFields <> " WHERE hash = ?") (Only hashText)

nameToTerm :: Connection -> Text -> IO (Maybe StoredTerm)
nameToTerm conn nameText =
  queryMaybeOne conn
    (selectStoredTermFields <> " WHERE (names = ? OR names LIKE ? OR names LIKE ? OR names LIKE ?) ORDER BY created_at DESC LIMIT 1")
    (nameText, nameText <> T.pack ",%", T.pack "%," <> nameText <> T.pack ",%", T.pack "%," <> nameText)

listStoredTerms :: Connection -> IO [StoredTerm]
listStoredTerms conn =
  query_ conn (selectStoredTermFields <> " ORDER BY created_at DESC")

storeEnvironment :: Connection -> Env -> IO [(String, Text)]
storeEnvironment conn env = do
  let defs = Map.toList $ Map.delete "!result" env
  let groupedDefs = Map.toList $ Map.fromListWith (++) [(term, [name]) | (name, term) <- defs]

  forM groupedDefs $ \(term, namesList) -> do
    hashVal <- storeTerm conn namesList term
    return (head namesList, hashVal)

loadTerm :: Connection -> String -> IO (Maybe T)
loadTerm conn identifier = do
  result <- getTerm conn (T.pack identifier)
  case result of
    Just storedTerm -> loadTree conn (termHash storedTerm)
    Nothing -> return Nothing

getTerm :: Connection -> Text -> IO (Maybe StoredTerm)
getTerm conn identifier = do
  if '#' `elem` (T.unpack identifier)
    then hashToTerm conn (T.pack $ drop 1 (T.unpack identifier))
    else nameToTerm conn identifier

loadEnvironment :: Connection -> IO Env
loadEnvironment conn = do
  terms <- listStoredTerms conn
  foldM addTermToEnv Map.empty terms
  where
    addTermToEnv env storedTerm = do
      maybeT <- loadTree conn (termHash storedTerm)
      case maybeT of
        Just t -> do
          let namesList = parseNameList (termNames storedTerm)
          return $ foldl (\e name -> Map.insert (T.unpack name) t e) env namesList
        Nothing -> return env

termVersions :: Connection -> String -> IO [(Text, T, Integer)]
termVersions conn name = do
  let nameText = T.pack name
  results <- query conn
    ("SELECT hash, created_at FROM terms WHERE (names = ? OR names LIKE ? OR names LIKE ? OR names LIKE ?) ORDER BY created_at DESC")
    (nameText, nameText <> T.pack ",%", T.pack "%," <> nameText <> T.pack ",%", T.pack "%," <> nameText)

  catMaybes <$> mapM (\(hashVal, timestamp) -> do
    maybeT <- loadTree conn hashVal
    return $ fmap (\t -> (hashVal, t, timestamp)) maybeT
    ) results

setTag :: Connection -> Text -> Text -> IO ()
setTag conn hash tagValue = do
  exists <- termExists conn hash
  if exists
    then do
      currentTagsQuery <- query conn "SELECT tags FROM terms WHERE hash = ?" (Only hash) :: IO [Only Text]
      case currentTagsQuery of
        [Only tagsText] -> do
          let tagsList = parseNameList tagsText
              newTagsList = tagValue : tagsList
              newTags = serializeNameList newTagsList
          execute conn "UPDATE terms SET tags = ? WHERE hash = ?" (newTags, hash)
        _ -> putStrLn $ "Term with hash " ++ T.unpack hash ++ " not found (should not happen if exists is true)"
    else
      putStrLn $ "Term with hash " ++ T.unpack hash ++ " does not exist"

termExists :: Connection -> Text -> IO Bool
termExists conn hash = do
  results <- query conn "SELECT 1 FROM terms WHERE hash = ? LIMIT 1" (Only hash) :: IO [[Int]]
  return $ not (null results)

termToTags :: Connection -> Text -> IO [Text]
termToTags conn hash = do
  tagsQuery <- query conn "SELECT tags FROM terms WHERE hash = ?" (Only hash) :: IO [Only Text]
  case tagsQuery of
    [Only tagsText] -> return $ parseNameList tagsText
    _ -> return []

tagToTerm :: Connection -> Text -> IO [StoredTerm]
tagToTerm conn tagValue = do
  let pattern = "%" <> tagValue <> "%"
  query conn (selectStoredTermFields <> " WHERE tags LIKE ? ORDER BY created_at DESC") (Only pattern)

allTermTags :: Connection -> IO [StoredTerm]
allTermTags conn = do
  query_ conn (selectStoredTermFields <> " WHERE tags IS NOT NULL AND tags != '' ORDER BY created_at DESC")

selectStoredTermFields :: Query
selectStoredTermFields = "SELECT hash, names, metadata, created_at, tags FROM terms"

queryMaybeOne :: (FromRow r, ToRow q) => Connection -> Query -> q -> IO (Maybe r)
queryMaybeOne conn qry params = do
  results <- query conn qry params
  case results of
    [row] -> return $ Just row
    _     -> return Nothing