module ContentStore where

import Research

import Control.Monad                  (foldM, forM_, void)
import Data.ByteString                (ByteString)
import Data.Char                      (isHexDigit)
import Data.List                      (nub, sort)
import Data.Maybe                     (catMaybes, fromMaybe)
import Data.Text                      (Text)
import Database.SQLite.Simple
import System.Directory               (createDirectoryIfMissing, getXdgDirectory, XdgDirectory(..))
import System.Environment             (lookupEnv)
import System.Exit                    (die)
import System.FilePath                ((</>), takeDirectory)

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
  setupDatabase conn
  return conn

-- | Initialise a database connection (file-backed or in-memory).
--   This is factored out so tests can reuse it with ":memory:".
setupDatabase :: Connection -> IO ()
setupDatabase conn = do
  execute_ conn "CREATE TABLE IF NOT EXISTS terms (\
                \hash TEXT PRIMARY KEY, \
                \names TEXT, \
                \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)"
  -- Seed canonical Leaf node payload (0x00)
  putMerkleNode conn NLeaf

-- | Create an in-memory ContentStore connection (for tests).
newContentStore :: IO Connection
newContentStore = do
  conn <- open ":memory:"
  setupDatabase conn
  return conn

getContentStorePath :: IO FilePath
getContentStorePath = do
  maybeLocalPath <- lookupEnv "TRICU_DB_PATH"
  case maybeLocalPath of
    Just p -> return p
    Nothing -> 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, metadata, tags) VALUES (?, ?, ?, ?)"
        (termHashText, allNamesToStore, 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)
    _ -> errorWithoutStackTrace $ "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 :: Connection -> MerkleHash -> IO (Maybe T)
loadTree conn h = do
  maybeNode <- getNodeMerkle conn h
  case maybeNode of
    Nothing -> return Nothing
    Just node -> Just <$> buildTree node
  where
    buildTree :: Node -> IO T
    buildTree NLeaf = return Leaf
    buildTree (NStem childHash) = do
      child <- fromMaybe (errorWithoutStackTrace "BUG: stored hash not found") <$> loadTree conn childHash
      return (Stem child)
    buildTree (NFork lHash rHash) = do
      left <- fromMaybe (errorWithoutStackTrace "BUG: stored hash not found") <$> loadTree conn lHash
      right <- fromMaybe (errorWithoutStackTrace "BUG: stored hash not found") <$> 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 conn Leaf = do
  putMerkleNode conn NLeaf
  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 ()
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) -> case namesList of
    _:_ -> void $ storeTerm conn namesList term
    _ -> errorWithoutStackTrace "storeEnvironment: empty names list"

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

-- | Resolve a user-supplied identifier (full/prefix hash, term name) to
--   a single term hash and the list of names bound to it.  Dies on
--   ambiguity or missing term (matching the CLI @export@ semantics).
resolveExportTarget :: Connection -> String -> IO (Text, [Text])
resolveExportTarget conn input = do
  let raw = T.pack $ dropWhile (== '#') input
  byName <- query conn
    "SELECT hash FROM terms WHERE (names = ? OR names LIKE ? OR names LIKE ? OR names LIKE ?) ORDER BY created_at DESC"
    (raw, raw <> T.pack ",%", T.pack "," <> raw <> T.pack ",%", T.pack "%," <> raw) :: IO [Only T.Text]
  case byName of
    [Only fullHash] -> namesForHash conn fullHash >>= \names -> return (fullHash, names)
    (_:_) -> die $ "Ambiguous term name: " ++ input
    [] -> do
      byHash <- query conn "SELECT hash FROM terms WHERE hash LIKE ? ORDER BY created_at DESC"
        (Only (raw <> T.pack "%")) :: IO [Only T.Text]
      case byHash of
        [Only fullHash] -> namesForHash conn fullHash >>= \names -> return (fullHash, names)
        [] -> if looksLikeHash raw
          then return (raw, [])
          else die $ "No term found matching: " ++ input
        _ -> die $ "Ambiguous hash prefix: " ++ input

namesForHash :: Connection -> Text -> IO [Text]
namesForHash conn h = do
  stored <- hashToTerm conn h
  return $ maybe [] (parseNameList . termNames) stored

-- | Return 'True' when @t@ looks like a full or partial SHA-256 hex hash.
looksLikeHash :: Text -> Bool
looksLikeHash t =
  let len = T.length t
  in len >= 16 && len <= 64 && T.all isHexDigit t