module ChuSQL.Algebra.Eval (evalRelOp, evalRelOpM) where

import ChuSQL.Algebra.Expr (evalCondForRow)
import ChuSQL.Algebra.Op
import ChuSQL.Algebra.Sort (sortRows)
import ChuSQL.Model
import ChuSQL.Storage (IndexResult (..), MonadStorage (..))
import ChuSQL.Syntax.AST (Expr (..))
import Control.Monad (filterM, foldM)
import qualified Data.HashMap.Strict as HM

-- 执行：按算子树算出结果行（Scan 阶段会加别名前缀）。

-- | 只保留清单里的列
project :: [String] -> Row -> Row
project ["*"] row = row
project cs row = [(c, v) | (c, v) <- row, c `elem` cs]

-- | 给列名加别名前缀
addPrefix :: Maybe String -> Row -> Row
addPrefix mAlias = map (\(k, v) -> (qualify mAlias k, v))

-- | 扫描一张表
evalScan :: Database -> Maybe String -> String -> Either String [Row]
evalScan db mAlias tbl = do
    table <- lookupTable db tbl
    Right (map (addPrefix mAlias) (tableRows table))

-- | 点查退化成扫描时用的条件
lookupCondition :: Maybe String -> String -> Int -> Expr
lookupCondition mAlias col k = Eq (Col (qualify mAlias col)) (LitInt k)

data EquiKeys = EquiKeys
    { ekLeft :: String
    , ekRight :: String
    }

-- | 一批行的列名（空批给空清单，此时一律退回嵌套循环）
rowsCols :: [Row] -> [String]
rowsCols (r : _) = map fst r
rowsCols [] = []

-- | 认出 a = b 这种两列等值条件，并且能分清谁在左谁在右
equiKeys :: [Row] -> [Row] -> Expr -> Maybe EquiKeys
equiKeys lrows rrows (Eq (Col a) (Col b))
    | a /= b, a `elem` lcols, b `elem` rcols = Just (EquiKeys a b)
    | a /= b, b `elem` lcols, a `elem` rcols = Just (EquiKeys b a)
  where
    lcols = rowsCols lrows
    rcols = rowsCols rrows
equiKeys _ _ _ = Nothing

-- | 哈希连接，输出顺序同嵌套循环
hashJoin :: EquiKeys -> [Row] -> [Row] -> [Row]
hashJoin keys lrows rrows = concatMap probe lrows
  where
    buckets =
        HM.fromListWith
            (flip (++))
            [(v, [r]) | r <- rrows, Just v <- [lookup (ekRight keys) r]]
    probe l = case lookup (ekLeft keys) l of
        Nothing -> []
        Just v -> [l ++ r | r <- HM.lookupDefault [] v buckets]

-- | 把两边的行按条件配对（边配边筛，不先建整张积表）
filterPairs :: (Row -> Either String Bool) -> [Row] -> [Row] -> Either String [Row]
filterPairs cond lrows rrows = go lrows []
  where
    go [] acc = Right (reverse acc)
    go (l : ls) acc = do
        acc' <- foldM (keep l) acc rrows
        go ls acc'
    keep l acc r = do
        let row = l ++ r
        ok <- cond row
        pure (if ok then row : acc else acc)

-- | 一次连接：能认出等值键就走哈希，否则退回嵌套循环
joinRows :: Expr -> [Row] -> [Row] -> Either String [Row]
joinRows cond lrows rrows = case equiKeys lrows rrows cond of
    Just keys -> Right (hashJoin keys lrows rrows)
    Nothing -> filterPairs (evalCondForRow cond) lrows rrows

-- | 纯求值（不需要存储）
evalRelOp :: Database -> RelOp -> Either String [Row]
evalRelOp db (Scan mAlias tbl) = evalScan db mAlias tbl
evalRelOp db (Lookup mAlias tbl col k) = do
    rows <- evalScan db mAlias tbl
    filterM (evalCondForRow (lookupCondition mAlias col k)) rows
evalRelOp db (Filter e op) = do
    rows <- evalRelOp db op
    filterM (evalCondForRow e) rows
evalRelOp db (Project cols op) = do
    rows <- evalRelOp db op
    Right (map (project cols) rows)
evalRelOp db (Sort spec op) = do
    rows <- evalRelOp db op
    Right (sortRows spec rows)
evalRelOp db (Limit n op) = do
    rows <- evalRelOp db op
    Right (take n rows)
evalRelOp db (Join left right cond) = do
    lrows <- evalRelOp db left
    rrows <- evalRelOp db right
    joinRows cond lrows rrows

-- | 单子求值：点查先问存储有无索引
evalRelOpM :: (MonadStorage m) => RelOp -> m (Either String [Row])
evalRelOpM (Scan mAlias t) = do
    result <- scan t
    pure $ do
        rows <- result
        case mAlias of
            Nothing -> Right rows
            Just _ -> Right (map (addPrefix mAlias) rows)
evalRelOpM (Lookup mAlias t col k) = do
    result <- lookupByColumn t col k
    case result of
        Left e -> pure (Left e)
        Right NoIndex -> do
            rows <- evalRelOpM (Scan mAlias t)
            pure $ do
                rs <- rows
                filterM (evalCondForRow (lookupCondition mAlias col k)) rs
        Right (IndexRow Nothing) -> pure (Right [])
        Right (IndexRow (Just row)) -> pure (Right [addPrefix mAlias row])
evalRelOpM (Filter p x) = do
    result <- evalRelOpM x
    pure $ do
        rows <- result
        filterM (evalCondForRow p) rows
evalRelOpM (Project cols x) = do
    result <- evalRelOpM x
    pure (fmap (map (project cols)) result)
evalRelOpM (Sort spec x) = do
    result <- evalRelOpM x
    pure (fmap (sortRows spec) result)
evalRelOpM (Limit n x) = do
    result <- evalRelOpM x
    pure (fmap (take n) result)
evalRelOpM (Join l r cond) = do
    ls <- evalRelOpM l
    rs <- evalRelOpM r
    pure $ do
        left <- ls
        right <- rs
        joinRows cond left right