{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveShow #-}
{-# LANGUAGE OverloadedStrings #-}

{-|
Module      : Assertica.Core.Termination
Description : Termination checking for recursive functions
Copyright   : (c) 2026 Ahmad Ali Parr
License     : BSL-1.0

This module implements termination checking to ensure all recursive definitions
are safe and will terminate. A function is considered terminating if:

1. **Direct Recursion**: All recursive calls are on structural subterms
   Example: `length (Cons h t) = 1 + length t` (t is subterm of Cons h t)

2. **Structural Recursion**: Recursion only happens on constructor arguments
   Example: `fib n = if n <= 1 then n else fib(n-1) + fib(n-2)` uses comparison guard

3. **Mutual Recursion**: Functions can call each other if there's a decreasing measure

4. **No Divergence**: Detect cycles without progress (e.g., f calls g calls f with same args)

DESIGN PRINCIPLES:
1. **Fail-closed**: Unsure recursion is rejected (no escape hatch)
2. **Structural focus**: Primary interest in structural recursion patterns
3. **Clear diagnostics**: Explain exactly why recursion is unsafe
4. **Integration with Type Checker**: Type environment needed for context

ALGORITHM:
1. Identify recursive calls in a function body
2. For each recursive call, verify:
   - It is called on a structural subterm of a constructor argument
   - The argument that changes is demonstrated to decrease
3. Build a call graph for mutual recursion
4. Verify all cycles have a decreasing component
-}

module Assertica.Core.Termination
  ( -- * Main termination checking
    checkTermination
  , checkTerminationMulti
  , needsTerminationCheck

    -- * Error type
  , TermError(..)
  , prettyTermError

    -- * Analysis utilities (for testing)
  , analyzeRecursiveCalls
  , RecursiveCall(..)
  , CallPattern(..)
  ) where

import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Set as Set
import Data.Set (Set)
import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map)
import Data.List (intercalate, nub)
import GHC.Generics (Generic)

import Assertica.Core.AST
import Assertica.Core.TypeEnv

-- ============================================================================
-- ERROR TYPES
-- ============================================================================

{-| Termination checking errors. -}
data TermError
  = NonTerminatingRecursion String
    -- ^ Explanation of why recursion doesn't terminate
  | ArgumentDoesNotDecrease String
    -- ^ Argument does not decrease on recursive call
  | UnsupportedRecursionPattern String
    -- ^ Recursion pattern not recognized as terminating
  | InfiniteCallCycle [Text]
    -- ^ Detected infinite call cycle (e.g., f -> g -> f)
  | MissingDecreasingMeasure String
    -- ^ Recursive call lacks proof of decrease
  deriving (Eq, Show, Generic)

{-| Pretty-print a termination error for user display. -}
prettyTermError :: TermError -> String
prettyTermError = \case
  NonTerminatingRecursion msg ->
    "Function does not terminate: " ++ msg
  ArgumentDoesNotDecrease msg ->
    "Recursive argument does not decrease: " ++ msg
  UnsupportedRecursionPattern msg ->
    "Recursion pattern not recognized as terminating: " ++ msg
  InfiniteCallCycle fnames ->
    "Infinite call cycle detected: " ++
    intercalate " -> " (map T.unpack fnames)
  MissingDecreasingMeasure msg ->
    "Missing proof that argument decreases: " ++ msg

-- ============================================================================
-- RECURSIVE CALL ANALYSIS
-- ============================================================================

{-| Information about a recursive call pattern. -}
data CallPattern
  = StructuralCall Text        -- ^ Call on subterm (e.g., "t" in Cons h t)
  | GuardedCall Text           -- ^ Call guarded by condition
  | UnknownPattern             -- ^ Unknown/unsupported pattern
  deriving (Eq, Show, Generic)

{-| A recursive call found in a function body. -}
data RecursiveCall = RecursiveCall
  { callName     :: Text           -- ^ Name of function being called
  , callPattern  :: CallPattern    -- ^ How the call is made
  , callTermStr  :: String         -- ^ String repr of what's passed
  }
  deriving (Eq, Show, Generic)

-- ============================================================================
-- TERMINATION CHECKING INTERFACE
-- ============================================================================

{-| Check if a term requires termination checking.
A term needs termination checking if it is a function definition with recursion.
-}
needsTerminationCheck :: TypeEnv -> Term -> Bool
needsTerminationCheck _env term = hasRecursion term

{-| Check termination of a single recursive function.

The function analyzes the term to ensure:
1. All recursive calls are on structural subterms
2. Each recursive call decreases a well-founded measure
3. No infinite loops without progress

Returns:
- Left TermError if termination cannot be verified
- Right () if the function is terminating
-}
checkTermination :: TypeEnv -> Text -> Term -> Either TermError ()
checkTermination _env funcName term = do
  if not (hasRecursion term)
    then Right ()
    else do
      calls <- analyzeRecursiveCalls funcName term
      case calls of
        [] -> Right ()  -- No recursive calls found
        _  -> verifyCallsTerminate funcName calls

{-| Check termination for multiple mutually recursive functions.

For mutual recursion, we build a call graph and verify that:
1. All cycles in the graph have a decreasing measure
2. No function calls itself without decrease
-}
checkTerminationMulti :: TypeEnv -> [(Text, Term)] -> Either TermError ()
checkTerminationMulti _env functions = do
  -- Check each function individually first
  mapM_ (\(name, term) -> checkTermination emptyEnv name term) functions

  -- For mutual recursion, analyze call graph (simplified)
  let callGraph = buildCallGraph functions
  detectInfiniteCycles callGraph

-- ============================================================================
-- RECURSIVE CALL DETECTION
-- ============================================================================

{-| Check if a term contains any recursive calls (references to itself or other functions). -}
hasRecursion :: Term -> Bool
hasRecursion = go
  where
    go = \case
      Var _ -> False
      Const _ -> False
      Lam _binder body -> go body
      App f x -> go f || go x
      Constr _name args -> any go args
      Case scrutinee clauses def ->
        go scrutinee || any (goClause) clauses || maybe False go def
      Let _binder e1 e2 -> go e1 || go e2
      Ann e _ty -> go e
      Forall _binder body -> go body
      ProofTerm _p -> False  -- Don't analyze proofs here
      Prop _prop -> False    -- Don't analyze propositions here

    goClause (Clause _pat body) = go body

{-| Analyze recursive calls in a term body.
Returns a list of recursive calls found, or an error if analysis fails.
-}
analyzeRecursiveCalls :: Text -> Term -> Either TermError [RecursiveCall]
analyzeRecursiveCalls funcName term =
  Right $ findCalls funcName term

{-| Find all calls to a specific function in a term. -}
findCalls :: Text -> Term -> [RecursiveCall]
findCalls targetFunc = go
  where
    go = \case
      Var _ -> []
      Const _ -> []
      Lam _binder body -> go body
      App f x ->
        let f_calls = go f
            x_calls = go x
            maybe_recursive = case f of
              Var (Var name _) | name == targetFunc ->
                [RecursiveCall targetFunc (analyzeCallArgument x) (prettyTerm x)]
              _ -> []
        in maybe_recursive ++ f_calls ++ x_calls
      Constr _name args -> concatMap go args
      Case scrutinee clauses def ->
        go scrutinee ++
        concatMap goClause clauses ++
        maybe [] go def
      Let _binder e1 e2 -> go e1 ++ go e2
      Ann e _ty -> go e
      Forall _binder body -> go body
      ProofTerm _p -> []
      Prop _prop -> []

    goClause (Clause _pat body) = go body

{-| Analyze the argument to a recursive call to determine the call pattern. -}
analyzeCallArgument :: Term -> CallPattern
analyzeCallArgument = \case
  -- If it's a pattern variable directly, it's structural if from a constructor arg
  Var _ -> StructuralCall "arg"

  -- If it's extracting a component from a constructor, likely structural
  Constr _ _ -> StructuralCall "constr"

  -- Function applications that decrease size (heuristic)
  App (Const (PrimOp _)) arg ->
    StructuralCall "computed"

  -- Default: unknown
  _ -> UnknownPattern

{-| Verify that recursive calls satisfy termination. -}
verifyCallsTerminate :: Text -> [RecursiveCall] -> Either TermError ()
verifyCallsTerminate _funcName calls = do
  -- For now, use a simple heuristic:
  -- If all calls use StructuralCall pattern, likely terminating
  let hasUnknown = any (\c -> callPattern c == UnknownPattern) calls

  if hasUnknown
    then do
      let unknownCall = head [c | c <- calls, callPattern c == UnknownPattern]
      Left $ UnsupportedRecursionPattern
        ("Recursion pattern in call to " ++ T.unpack (callName unknownCall) ++
         " with argument " ++ callTermStr unknownCall ++
         " is not recognized as terminating. " ++
         "Only structural recursion on constructor arguments is supported.")
    else Right ()

-- ============================================================================
-- CALL GRAPH ANALYSIS (for mutual recursion)
-- ============================================================================

{-| Build a call graph from a set of mutually recursive functions. -}
buildCallGraph :: [(Text, Term)] -> Map Text [Text]
buildCallGraph functions =
  Map.fromList [(name, findDirectCalls term) | (name, term) <- functions]

{-| Find all direct function calls in a term (by name). -}
findDirectCalls :: Term -> [Text]
findDirectCalls = nub . go
  where
    go = \case
      Var (Var name _) -> [name]
      Const _ -> []
      Lam _binder body -> go body
      App f x -> go f ++ go x
      Constr _name args -> concatMap go args
      Case scrutinee clauses def ->
        go scrutinee ++ concatMap goClause clauses ++ maybe [] go def
      Let _binder e1 e2 -> go e1 ++ go e2
      Ann e _ty -> go e
      Forall _binder body -> go body
      ProofTerm _p -> []
      Prop _prop -> []

    goClause (Clause _pat body) = go body

{-| Detect infinite call cycles in a call graph. -}
detectInfiniteCycles :: Map Text [Text] -> Either TermError ()
detectInfiniteCycles graph = do
  case findCycle graph of
    Nothing -> Right ()
    Just cycle_path ->
      Left $ InfiniteCallCycle cycle_path

{-| Find a cycle in a call graph (simplified DFS). -}
findCycle :: Map Text [Text] -> Maybe [Text]
findCycle graph =
  case Map.keys graph of
    [] -> Nothing
    (start:_) -> dfs graph start start Set.empty

{-| Depth-first search to find a cycle. -}
dfs :: Map Text [Text] -> Text -> Text -> Set Text -> Maybe [Text]
dfs graph current start visited =
  if Set.member current visited
    then if current == start && not (Set.null visited)
         then Just [current]
         else Nothing
    else case Map.lookup current graph of
      Nothing -> Nothing
      Just callees ->
        case concat [maybe [] (\path -> [current] ++ path) (dfs graph callee start (Set.insert current visited)) | callee <- callees] of
          [] -> Nothing
          path -> Just path

-- ============================================================================
-- UTILITIES
-- ============================================================================

{-| Helper: check if a term decreases relative to an input argument.
This is a heuristic check used for simple termination patterns.
-}
decreasesOn :: Term -> Term -> Bool
decreasesOn _recursiveArg _inputArg =
  -- Simplified: in a real implementation, we'd check structural subterm relationship
  True

{-| Pretty-print a term (reuse from AST). -}
prettyTerm :: Term -> String
prettyTerm = Assertica.Core.AST.prettyTerm