## CDCL Solver for Dependency Resolution
##
## This module implements a Conflict-Driven Clause Learning (CDCL) SAT solver
## adapted for package dependency resolution with the PubGrub algorithm.
##
## Philosophy:
## - Start with root requirements (unit clauses)
## - Make decisions (select package versions)
## - Propagate implications (unit propagation)
## - Detect conflicts
## - Learn from conflicts (add new clauses)
## - Backjump to root cause (non-chronological backtracking)
##
## Key Concepts:
## - Decision: Choosing a package version to install
## - Implication: Forced choice due to unit propagation
## - Conflict: Incompatible assignments detected
## - Learned Clause: New constraint derived from conflict analysis
## - Backjumping: Jump to earliest decision causing conflict

import std/[tables, sets, options, sequtils, algorithm]
import ./cnf_translator
import ./solver_types
import ./variant_types
import ../manifest_parser

type
  ## Assignment type (decision vs implication)
  AssignmentType* = enum
    Decision,      ## User choice or heuristic selection
    Implication    ## Forced by unit propagation
  
  ## A variable assignment in the solver
  SolverAssignment* = object
    variable*: BoolVar
    value*: bool  ## true = selected, false = not selected
    assignmentType*: AssignmentType
    decisionLevel*: int
    antecedent*: Option[Clause]  ## The clause that forced this (for implications)
  
  ## Conflict information
  Conflict* = object
    clause*: Clause
    assignments*: seq[SolverAssignment]
  
  ## The CDCL solver state
  CDCLSolver* = object
    formula*: CNFFormula
    assignments*: Table[BoolVar, SolverAssignment]
    decisionLevel*: int
    learnedClauses*: seq[Clause]
    propagationQueue*: seq[BoolVar]
  
  ## Solver result
  SolverResult* = object
    case isSat*: bool
    of true:
      model*: Table[BoolVar, bool]
    of false:
      conflict*: Conflict

# --- Assignment Operations ---

proc isAssigned*(solver: CDCLSolver, variable: BoolVar): bool =
  ## Check if a variable has been assigned
  result = solver.assignments.hasKey(variable)

proc getAssignment*(solver: CDCLSolver, variable: BoolVar): Option[SolverAssignment] =
  ## Get the assignment for a variable
  if solver.assignments.hasKey(variable):
    return some(solver.assignments[variable])
  else:
    return none(SolverAssignment)

proc getValue*(solver: CDCLSolver, variable: BoolVar): Option[bool] =
  ## Get the value of a variable
  if solver.assignments.hasKey(variable):
    return some(solver.assignments[variable].value)
  else:
    return none(bool)

proc assign*(solver: var CDCLSolver, variable: BoolVar, value: bool, 
             assignmentType: AssignmentType, antecedent: Option[Clause] = none(Clause)) =
  ## Assign a value to a variable
  solver.assignments[variable] = SolverAssignment(
    variable: variable,
    value: value,
    assignmentType: assignmentType,
    decisionLevel: solver.decisionLevel,
    antecedent: antecedent
  )
  
  # Add to propagation queue if this is a decision
  if assignmentType == Decision:
    solver.propagationQueue.add(variable)

proc unassign*(solver: var CDCLSolver, variable: BoolVar) =
  ## Remove an assignment
  solver.assignments.del(variable)

# --- Clause Evaluation ---

proc evaluateLiteral*(solver: CDCLSolver, literal: Literal): Option[bool] =
  ## Evaluate a literal given current assignments
  ## Returns: Some(true) if satisfied, Some(false) if falsified, None if unassigned
  
  let varValue = solver.getValue(literal.variable)
  if varValue.isNone:
    return none(bool)
  
  let value = varValue.get()
  if literal.isNegated:
    return some(not value)
  else:
    return some(value)

proc evaluateClause*(solver: CDCLSolver, clause: Clause): Option[bool] =
  ## Evaluate a clause given current assignments
  ## Returns: Some(true) if satisfied, Some(false) if falsified, None if undetermined
  
  var hasUnassigned = false
  
  for literal in clause.literals:
    let litValue = solver.evaluateLiteral(literal)
    
    if litValue.isSome:
      if litValue.get():
        # Clause is satisfied (at least one literal is true)
        return some(true)
    else:
      hasUnassigned = true
  
  if hasUnassigned:
    # Clause is undetermined (has unassigned literals)
    return none(bool)
  else:
    # All literals are false, clause is falsified
    return some(false)

proc isUnitClause*(solver: CDCLSolver, clause: Clause): Option[Literal] =
  ## Check if a clause is unit (exactly one unassigned literal, rest false)
  ## Returns the unassigned literal if unit, None otherwise
  
  var unassignedLiteral: Option[Literal] = none(Literal)
  var unassignedCount = 0
  
  for literal in clause.literals:
    let litValue = solver.evaluateLiteral(literal)
    
    if litValue.isNone:
      # Unassigned literal
      unassignedCount += 1
      unassignedLiteral = some(literal)
      if unassignedCount > 1:
        return none(Literal)  # More than one unassigned
    elif litValue.get():
      # Literal is true, clause is satisfied
      return none(Literal)
  
  if unassignedCount == 1:
    return unassignedLiteral
  else:
    return none(Literal)

# --- Unit Propagation ---

proc unitPropagate*(solver: var CDCLSolver): Option[Conflict] =
  ## Perform unit propagation (Boolean Constraint Propagation)
  ## Returns a conflict if one is detected, None otherwise
  ##
  ## Requirements: 5.1 - Use PubGrub algorithm with CDCL
  
  var changed = true
  while changed:
    changed = false
    
    # Check all clauses for unit clauses
    for clause in solver.formula.clauses:
      let clauseValue = solver.evaluateClause(clause)
      
      if clauseValue.isSome and not clauseValue.get():
        # Clause is falsified - conflict!
        return some(Conflict(
          clause: clause,
          assignments: solver.assignments.values.toSeq
        ))
      
      let unitLit = solver.isUnitClause(clause)
      if unitLit.isSome:
        let lit = unitLit.get()
        
        # Check if already assigned
        if solver.isAssigned(lit.variable):
          let currentValue = solver.getValue(lit.variable).get()
          let requiredValue = not lit.isNegated
          
          if currentValue != requiredValue:
            # Conflict: variable must be both true and false
            return some(Conflict(
              clause: clause,
              assignments: solver.assignments.values.toSeq
            ))
        else:
          # Assign the variable to satisfy the unit clause
          let value = not lit.isNegated
          solver.assign(lit.variable, value, Implication, some(clause))
          changed = true
    
    # Check learned clauses too
    for clause in solver.learnedClauses:
      let clauseValue = solver.evaluateClause(clause)
      
      if clauseValue.isSome and not clauseValue.get():
        # Clause is falsified - conflict!
        return some(Conflict(
          clause: clause,
          assignments: solver.assignments.values.toSeq
        ))
      
      let unitLit = solver.isUnitClause(clause)
      if unitLit.isSome:
        let lit = unitLit.get()
        
        if not solver.isAssigned(lit.variable):
          let value = not lit.isNegated
          solver.assign(lit.variable, value, Implication, some(clause))
          changed = true
  
  return none(Conflict)

# --- Decision Heuristics ---

proc selectUnassignedVariable*(solver: CDCLSolver): Option[BoolVar] =
  ## Select an unassigned variable using a heuristic
  ## For now, we use a simple first-unassigned heuristic
  ## TODO: Implement VSIDS or other advanced heuristics
  
  for variable, _ in solver.formula.variables.pairs:
    if not solver.isAssigned(variable):
      return some(variable)
  
  return none(BoolVar)

# --- Conflict Analysis ---

proc analyzeConflict*(solver: CDCLSolver, conflict: Conflict): Clause =
  ## Analyze a conflict and learn a new clause
  ## This implements the "first UIP" (Unique Implication Point) scheme
  ##
  ## Requirements: 5.2 - Learn new incompatibility clause from conflicts
  
  # 1. Initialize resolution
  # Start with the conflict clause
  var currentClauseLiterals = conflict.clause.literals
  
  # We want to resolve literals that were assigned at the current decision level
  # until only one remains (the UIP).
  
  # For this MVP, we'll stick to a simpler "block this assignment" strategy
  # but with a bit more intelligence: we'll include the decision variables
  # that led to this conflict.
  
  var learnedLiterals: seq[Literal] = @[]
  var seenVariables = initHashSet[BoolVar]()
  
  # Collect all decision variables that are antecedents of the conflict
  for assignment in conflict.assignments:
    if assignment.assignmentType == Decision:
      if assignment.variable notin seenVariables:
        seenVariables.incl(assignment.variable)
        # Negate the decision
        learnedLiterals.add(makeLiteral(assignment.variable, isNegated = not assignment.value))
  
  # If we found decisions, use them. Otherwise fall back to the conflict clause.
  if learnedLiterals.len > 0:
    return makeClause(learnedLiterals, reason = "Learned from conflict decision path")
  else:
    return conflict.clause

proc findBackjumpLevel*(solver: CDCLSolver, learnedClause: Clause): int =
  ## Find the decision level to backjump to
  ## This is the second-highest decision level in the learned clause
  ##
  ## Requirements: 5.3 - Backjump to earliest decision causing conflict
  
  var levels: seq[int] = @[]
  
  for literal in learnedClause.literals:
    if solver.isAssigned(literal.variable):
      let assignment = solver.getAssignment(literal.variable).get()
      if assignment.decisionLevel notin levels:
        levels.add(assignment.decisionLevel)
  
  if levels.len == 0:
    return 0
  
  levels.sort()
  
  if levels.len == 1:
    return max(0, levels[0] - 1)
  else:
    # Return second-highest level
    return levels[levels.len - 2]

proc backjump*(solver: var CDCLSolver, level: int) =
  ## Backjump to a specific decision level
  ## Remove all assignments made after that level
  ##
  ## Requirements: 5.3 - Backjump to earliest decision causing conflict
  
  var toRemove: seq[BoolVar] = @[]
  
  for variable, assignment in solver.assignments.pairs:
    if assignment.decisionLevel > level:
      toRemove.add(variable)
  
  for variable in toRemove:
    solver.unassign(variable)
  
  solver.decisionLevel = level
  solver.propagationQueue = @[]

# --- Main Solver Loop ---

proc solve*(solver: var CDCLSolver): SolverResult =
  ## Main CDCL solving loop
  ## Returns SAT with model if satisfiable, UNSAT with conflict if not
  ##
  ## Requirements: 5.1, 5.2, 5.3, 5.4, 5.5
  
  # Initial unit propagation
  let initialConflict = solver.unitPropagate()
  if initialConflict.isSome:
    # Formula is unsatisfiable at decision level 0
    return SolverResult(isSat: false, conflict: initialConflict.get())
  
  # Main CDCL loop
  while true:
    # Check if all variables are assigned
    let unassignedVar = solver.selectUnassignedVariable()
    
    if unassignedVar.isNone:
      # All variables assigned, formula is satisfied!
      var model = initTable[BoolVar, bool]()
      for variable, assignment in solver.assignments.pairs:
        model[variable] = assignment.value
      return SolverResult(isSat: true, model: model)
    
    # Make a decision
    solver.decisionLevel += 1
    let variable = unassignedVar.get()
    solver.assign(variable, true, Decision)  # Try true first
    
    # Propagate implications
    let conflict = solver.unitPropagate()
    
    if conflict.isSome:
      # Conflict detected!
      if solver.decisionLevel == 0:
        # Conflict at decision level 0 - unsatisfiable
        return SolverResult(isSat: false, conflict: conflict.get())
      
      # Analyze conflict and learn
      let learnedClause = solver.analyzeConflict(conflict.get())
      solver.learnedClauses.add(learnedClause)
      
      # Backjump
      let backjumpLevel = solver.findBackjumpLevel(learnedClause)
      solver.backjump(backjumpLevel)

# --- Solver Construction ---

proc newCDCLSolver*(formula: CNFFormula): CDCLSolver =
  ## Create a new CDCL solver for a CNF formula
  result = CDCLSolver(
    formula: formula,
    assignments: initTable[BoolVar, SolverAssignment](),
    decisionLevel: 0,
    learnedClauses: @[],
    propagationQueue: @[]
  )

# --- String Representations ---

proc `$`*(assignment: SolverAssignment): string =
  ## String representation of an assignment
  result = $assignment.variable & " = " & $assignment.value
  result.add(" @" & $assignment.decisionLevel)
  if assignment.assignmentType == Decision:
    result.add(" (decision)")
  else:
    result.add(" (implied)")

proc `$`*(conflict: Conflict): string =
  ## String representation of a conflict
  result = "Conflict in clause: " & $conflict.clause

proc `$`*(solverResult: SolverResult): string =
  ## String representation of solver result
  if solverResult.isSat:
    result = "SAT (" & $solverResult.model.len & " variables assigned)"
  else:
    result = "UNSAT: " & $solverResult.conflict
