module Eval where

import ContentStore
import Parser
import Research

import Control.Monad (forM_, foldM)
import Data.List (partition, (\\))
import Data.Map  (Map)
import Database.SQLite.Simple
import qualified Data.Map as Map
import qualified Data.Set as Set
import qualified Data.Text as T
import Data.List (foldl')

evalSingle :: Env -> TricuAST -> Env
evalSingle env term
  | SDef name [] body <- term
  = case Map.lookup name env of
      Just existingValue
        | existingValue == evalASTSync env body -> env
        | otherwise
        -> let res = evalASTSync env body
           in Map.insert "!result" res (Map.insert name res env)
      Nothing
        -> let res = evalASTSync env body
           in Map.insert "!result" res (Map.insert name res env)
  | SApp func arg <- term
  = let res = apply (evalASTSync env func) (evalASTSync env arg)
    in Map.insert "!result" res env
  | SVar name Nothing <- term
  = case Map.lookup name env of
      Just v -> Map.insert "!result" v env
      Nothing -> errorWithoutStackTrace $ "Variable " ++ name ++ " not defined"
  | SVar name (Just hash) <- term
  = errorWithoutStackTrace $ "Hash-specific variable lookup not supported in local evaluation: " ++ name ++ "#" ++ hash
  | otherwise
  = let res = evalASTSync env term
    in Map.insert "!result" res env

evalTricu :: Env -> [TricuAST] -> Env
evalTricu env x = go env (reorderDefs env x)
  where
    go env []     = env
    go env [x]    =
      let updatedEnv = evalSingle env x
      in Map.insert "!result" (result updatedEnv) updatedEnv
    go env (x:xs) =
      evalTricu (evalSingle env x) xs

evalASTSync :: Env -> TricuAST -> T
evalASTSync env term = case term of
  SLambda _ _ -> evalASTSync env (elimLambda term)
  SVar name Nothing -> case Map.lookup name env of
    Just v -> v
    Nothing -> errorWithoutStackTrace $ "Variable " ++ name ++ " not defined"
  SVar name (Just hash) ->
    case Map.lookup (name ++ "#" ++ hash) env of
      Just v -> v
      Nothing -> errorWithoutStackTrace $
        "Variable " ++ name ++ " with hash " ++ hash ++ " not found in environment"
  TLeaf -> Leaf
  TStem t -> Stem (evalASTSync env t)
  TFork t u -> Fork (evalASTSync env t) (evalASTSync env u)
  SApp t u -> apply (evalASTSync env t) (evalASTSync env u)
  SStr s -> ofString s
  SInt n -> ofNumber n
  SList xs -> ofList (map (evalASTSync env) xs)
  SEmpty -> Leaf
  _ -> errorWithoutStackTrace $ "Unexpected AST term: " ++ show term

evalAST :: Maybe Connection -> Map.Map String T.Text -> TricuAST -> IO T
evalAST mconn selectedVersions ast = do
  let varNames = collectVarNames ast
  resolvedEnv <- resolveTermsFromStore mconn selectedVersions varNames
  return $ evalASTSync resolvedEnv ast

collectVarNames :: TricuAST -> [(String, Maybe String)]
collectVarNames = go []
  where
    go acc (SVar name mhash) = (name, mhash) : acc
    go acc (SApp t u) = go (go acc t) u
    go acc (SLambda vars body) =
      let boundVars = Set.fromList vars
          collected = go [] body
      in acc ++ filter (\(name, _) -> not $ Set.member name boundVars) collected
    go acc (TStem t) = go acc t
    go acc (TFork t u) = go (go acc t) u
    go acc (SList xs) = foldl' go acc xs
    go acc _ = acc

resolveTermsFromStore :: Maybe Connection -> Map.Map String T.Text -> [(String, Maybe String)] -> IO Env
resolveTermsFromStore Nothing _ _ = return Map.empty
resolveTermsFromStore (Just conn) selectedVersions varNames = do
  foldM (\env (name, mhash) -> do
    term <- resolveTermFromStore conn selectedVersions name mhash
    case term of
      Just t -> return $ Map.insert (getVarKey name mhash) t env
      Nothing -> return env
      ) Map.empty varNames
  where
    getVarKey name Nothing = name
    getVarKey name (Just hash) = name ++ "#" ++ hash

resolveTermFromStore :: Connection -> Map.Map String T.Text -> String -> Maybe String -> IO (Maybe T)
resolveTermFromStore conn selectedVersions name mhash = case mhash of
  Just hashPrefix -> do
    versions <- termVersions conn name
    let matchingVersions = filter (\(hash, _, _) ->
          T.isPrefixOf (T.pack hashPrefix) hash) versions
    case matchingVersions of
      [] -> return Nothing
      [(_, term, _)] -> return $ Just term
      _ -> return Nothing -- Ambiguous or too many matches
  Nothing -> case Map.lookup name selectedVersions of
    Just hash -> loadTree conn hash
    Nothing -> do
      versions <- termVersions conn name
      case versions of
        [] -> return Nothing
        [(_, term, _)] -> return $ Just term
        _ -> return $ Just $ (\(_, t, _) -> t) $ head versions

elimLambda :: TricuAST -> TricuAST
elimLambda = go
  where
    go term
      | etaReduction term   = elimLambda $ etaReduceResult term
      | triagePattern term  = _TRI
      | composePattern term = _B
      | lambdaList term     = elimLambda $ lambdaListResult term
      | nestedLambda term   = nestedLambdaResult term
      | application term    = applicationResult term
      | isSList term        = slistTransform term
      | otherwise           = term

    etaReduction (SLambda [v] (SApp f (SVar x Nothing))) = v == x && not (isFree v f)
    etaReduction _ = False
    etaReduceResult (SLambda [_] (SApp f _)) = f

    triagePattern (SLambda [a] (SLambda [b] (SLambda [c] body))) = body == triageBody a b c
    triagePattern _ = False

    composePattern (SLambda [f] (SLambda [g] (SLambda [x] body))) = body == composeBody f g x
    composePattern _ = False

    lambdaList (SLambda [_] (SList _)) = True
    lambdaList _ = False
    lambdaListResult (SLambda [v] (SList xs)) = SLambda [v] (foldr wrapTLeaf TLeaf xs)
    wrapTLeaf m r = SApp (SApp TLeaf m) r

    nestedLambda (SLambda (_:_) _) = True
    nestedLambda _ = False
    nestedLambdaResult (SLambda (v:vs) body)
      | null vs   = toSKI v (go body) -- Changed elimLambda to go
      | otherwise = go (SLambda [v] (SLambda vs body)) -- Changed elimLambda to go

    application (SApp _ _) = True
    application _ = False
    applicationResult (SApp f g) = SApp (go f) (go g) -- Changed elimLambda to go

    isSList (SList _) = True
    isSList _         = False

    slistTransform :: TricuAST -> TricuAST
    slistTransform (SList xs) = foldr (\m r -> SApp (SApp TLeaf (go m)) r) TLeaf xs
    slistTransform ast = ast -- Should not be reached if isSList is the guard

    toSKI x (SVar y Nothing)
      | x == y = _I
      | otherwise = SApp _K (SVar y Nothing)
    toSKI x (SApp m n) = SApp (SApp _S (toSKI x m)) (toSKI x n)
    toSKI x (SLambda [y] body) = toSKI x (toSKI y body) -- This should ideally not happen if lambdas are fully eliminated first
    toSKI _ sl@(SList _) = SApp _K (go sl) -- Ensure SList itself is transformed if somehow passed to toSKI directly
    toSKI _ term = SApp _K term

    _S   = parseSingle "t (t (t t t)) t"
    _K   = parseSingle "t t"
    _I   = parseSingle "t (t (t t)) t"
    _B   = parseSingle "t (t (t t (t (t (t t t)) t))) (t t)"
    _TRI = parseSingle "t (t (t t (t (t (t t t))))) t"

    triageBody a b c = SApp (SApp TLeaf (SApp (SApp TLeaf (SVar a Nothing)) (SVar b Nothing))) (SVar c Nothing)
    composeBody f g x = SApp (SVar f Nothing) (SVar g Nothing) -- Note: This might not be the standard B combinator body f(g x)

isFree :: String -> TricuAST -> Bool
isFree x = Set.member x . freeVars

freeVars :: TricuAST -> Set.Set String
freeVars (SVar v Nothing) = Set.singleton v
freeVars (SVar v (Just _)) = Set.singleton v
freeVars (SApp t u) = Set.union (freeVars t) (freeVars u)
freeVars (SLambda vs body) = Set.difference (freeVars body) (Set.fromList vs)
freeVars _ = Set.empty

reorderDefs :: Env -> [TricuAST] -> [TricuAST]
reorderDefs env defs
  | not (null missingDeps) =
      errorWithoutStackTrace $
        "Missing dependencies detected: " ++ show missingDeps
  | otherwise = orderedDefs ++ others
  where
    (defsOnly, others) = partition isDef defs
    defNames = [ name | SDef name _ _ <- defsOnly ]

    defsWithFreeVars = [(def, freeVars body) | def@(SDef _ _ body) <- defsOnly]

    graph = buildDepGraph defsOnly
    sortedDefs = sortDeps graph
    defMap = Map.fromList [(name, def) | def@(SDef name _ _) <- defsOnly]
    orderedDefs = map (defMap Map.!) sortedDefs

    freeVarsDefs = foldMap snd defsWithFreeVars
    freeVarsOthers = foldMap freeVars others
    allFreeVars = freeVarsDefs <> freeVarsOthers
    validNames = Set.fromList defNames `Set.union` Set.fromList (Map.keys env)
    missingDeps = Set.toList (allFreeVars `Set.difference` validNames)

    isDef SDef {} = True
    isDef _       = False

buildDepGraph :: [TricuAST] -> Map.Map String (Set.Set String)
buildDepGraph topDefs
  | not (null conflictingDefs) =
      errorWithoutStackTrace $
        "Conflicting definitions detected: " ++ show conflictingDefs
  | otherwise =
      Map.fromList
        [ (name, depends topDefs (SDef name [] body))
        | SDef name _ body <- topDefs]
  where
    defsMap = Map.fromListWith (++)
      [(name, [(name, body)]) | SDef name _ body <- topDefs]

    conflictingDefs =
      [ name
      | (name, defs) <- Map.toList defsMap
      , let bodies = map snd defs
      , not $ all (== head bodies) (tail bodies)
      ]

sortDeps :: Map.Map String (Set.Set String) -> [String]
sortDeps graph = go [] Set.empty (Map.keys graph)
  where
    go sorted sortedSet [] = sorted
    go sorted sortedSet remaining =
      let ready = [ name | name <- remaining
                        , let deps = Map.findWithDefault Set.empty name graph
                        , Set.isSubsetOf deps sortedSet ]
          notReady = remaining \\ ready
      in if null ready
         then errorWithoutStackTrace
          "ERROR: Cyclic dependency detected and prohibited.\n\
          \RESOLVE: Use nested lambdas."
         else go (sorted ++ ready)
                 (Set.union sortedSet (Set.fromList ready))
                 notReady

depends :: [TricuAST] -> TricuAST -> Set.Set String
depends topDefs (SDef _ _ body) =
  Set.intersection
    (Set.fromList [n | SDef n _ _ <- topDefs])
    (freeVars body)
depends _ _ = Set.empty

result :: Env -> T
result r = case Map.lookup "!result" r of
  Just a -> a
  Nothing -> errorWithoutStackTrace "No !result field found in provided env"

mainResult :: Env -> T
mainResult r = case Map.lookup "main" r of
  Just  a -> a
  Nothing -> errorWithoutStackTrace "No valid definition for `main` found."

evalWithEnv :: Env -> Maybe Connection -> Map.Map String T.Text -> TricuAST -> IO T
evalWithEnv env mconn selectedVersions ast = do
  let varNames = findVarNames ast
  resolvedEnv <- case mconn of
    Just conn -> foldM (\e name ->
      if Map.member name e
        then return e
        else do
          mterm <- resolveTermFromStore conn selectedVersions name Nothing
          case mterm of
            Just term -> return $ Map.insert name term e
            Nothing -> return e
      ) env varNames
    Nothing -> return env
  return $ evalASTSync resolvedEnv ast

findVarNames :: TricuAST -> [String]
findVarNames ast = case ast of
  SVar name _ -> [name]
  SApp a b -> findVarNames a ++ findVarNames b
  SLambda args body -> findVarNames body \\ args
  SDef name args body -> name : (findVarNames body \\ args)
  _ -> []