package de.tum.workflows

import scala.util.parsing.combinator._
import de.tum.workflows.blocks._
import de.tum.workflows.foltl.FOLTL._
import de.tum.workflows.blocks.Workflow._
import com.typesafe.scalalogging.LazyLogging

object WorkflowParser extends RegexParsers with LazyLogging {
  
  def addChoicePredicates(blocks: List[Block]) = {
    def addChoice(blocks: List[Block], next:Int):List[Block] = {
      blocks match {
        case ForallMayBlock(v, _, stmts) :: xs => ForallMayBlock(v, s"choice$next", stmts) :: addChoice(xs, next + 1)
        case Loop(blocks) :: xs => Loop(addChoice(blocks, next)) :: addChoice(xs, next + blocks.size)
        case x :: xs => x :: addChoice(xs, next)
        case Nil => Nil
      }
    }
    addChoice(blocks, 0)
  }
  
  def allpredicates(list: List[Block]) = {
    def predicatesStmt(s: Statement):Set[Fun] = {
      // collect from guard
      val inguard = s.guard.collect({
        case f:Fun => List(f)
      })
      // changing function
      val ch = Fun(s.fun, s.tuple)
      inguard :+ ch toSet
    }
    
    def predicates(b: Block): Set[Fun] = {
      b match {
        case Loop(blocks) => blocks flatMap(predicates _) toSet
        case b:SimpleBlock => b.steps.flatMap(predicatesStmt _) toSet
        case NondetChoice(l,r) => (l.flatMap(predicates) ++ r.flatMap(predicates)) toSet
      }
    }
    
    list flatMap predicates
  }
  
  def checkSig(list: List[Block]) = {
    val preds = allpredicates(list)
    
    // Filter to just the first
    val grouped = preds groupBy (_.name)
    
    // Check arities
    val checks = for ((k, list) <- grouped) yield {
      val sizes = list map (_.params.size)
      // all equal
      (k, sizes.distinct.size == 1)
    }
    
    for ((k, valid) <- checks) {
      if (!valid) {
        logger.error(s"Predicate $k appears with conflicting arities")
      }
    }
    
    // all valid?
    checks.forall(x => x._2)
  }
  
  def buildSig(list: List[Block]) = {
    val preds = allpredicates(list)
    
    // Filter to just the first
    val filtered = preds groupBy (_.name) map (_._2.head)
    val (oracles, rels) = filtered.partition(_.name.startsWith("O"))
    Signature(oracles toSet, rels toSet)
  }
  
  def typeMap(typedvars: List[Var]) = {
    typedvars.map(v => v.name -> v.typ) toMap
  }
  
  def inferType(typemap: Map[String, String], v:Var) = {
    Var(v.name, typemap(v.name))
  }
  
  def inferTypes(typemap: Map[String, String], stmt: Statement) = {
    
    val newguard = stmt.guard.everywhere({
      case Fun(n, tr, params) => Fun(n, tr, params map (inferType(typemap, _))) 
    })
    val newtup = stmt.tuple.map (inferType(typemap, _))
    
    // TODO this could be improved
    stmt match {
      case s:Add => Add(newguard, stmt.fun, newtup)
      case s:Remove => Remove(newguard, stmt.fun, newtup)
    }
  }
  
  def validblock(bound: List[Var], stmts:List[Statement]) = {
    // no duplicate variables
    val distinct = bound.distinct.size == bound.size
    
    // all variable names contained in bound
    val names = bound.map(_.name)
    val unbound = stmts.flatMap(_.freeVars()).filterNot(v => names.contains(v.name))
    if (!unbound.isEmpty) {
      logger.error(s"Variables $unbound appear unbound.")
    }
    distinct && unbound.isEmpty
  }

  def IDENT = "[a-zA-Z][a-zA-Z0-9]*".r
  def VAR =  IDENT ^^ { str => Var(str) }
  def TYPEDVAR = IDENT ~ ((":" ~> IDENT)?) ^^ { case str ~ typ => if (typ.isDefined) Var(str, typ.get) else Var(str) } 
  
  // logics
  def tt = "True" ^^^ True
  def ff = "False" ^^^ False
  def AND = repsep(SIMPLETERM, "∧") ^^ { case ts => And.make(ts) } 
  def OR = repsep(SIMPLETERM, "∨") ^^ { case ts => Or.make(ts) } 
  def NEG = ("¬" | "!") ~ SIMPLETERM ^^ { case _ ~ t => Neg(t) } 
  
  // temporals
  def GLOBALLY = ("G" | "☐") ~> TERM ^^ { case t => Globally(t) }
  def FINALLY = ("F" | "♢") ~> TERM ^^ { case t => Finally(t) }
  
  def TUPLE = "(" ~> repsep(TYPEDVAR, ",") <~ ")" 
  def FUN = VAR ~ TUPLE ^^ { case f ~ tup => Fun(f.name, tup) }
  
  def SIMPLETERM = (tt | ff | FUN )
  def TERM:Parser[Term] = (SIMPLETERM | GLOBALLY | FINALLY | "(" ~> AND <~ ")" | "(" ~> OR <~ ")" | NEG )
  
  def ADD = TERM ~ ("→" | "->") ~ VAR ~ "+=" ~ TUPLE ^^ { case guard ~ _ ~ fun ~ _ ~ tup => Add(guard, fun.name, tup) }
  def REM = TERM ~ ("→" | "->") ~ VAR ~ "-=" ~ TUPLE ^^ { case guard ~ _ ~ fun ~ _ ~ tup => Remove(guard, fun.name, tup) }
  
  def STMT = ADD | REM
  
  // TODO: There could be several - one per Oracle
  def DECLASS = (FUN <~ ":") ~ TERM ^^ { case f ~ t => (f,t) } 
  
  def BLOCK = "forall" ~> ("may"?) ~ repsep(TYPEDVAR, ",") ~ (STMT+) ^? ({ 
    // all free variables contained
    case may ~ vars ~ stmts if validblock(vars, stmts) => {
      val typemap = typeMap(vars)
      val typedstmts = stmts map (inferTypes(typemap, _))
      if (may.isEmpty) {
        ForallBlock(vars, typedstmts)
      } else {
        ForallMayBlock(vars, "", typedstmts)
      }
    }
  }, t => s"$t is not a valid block (check free variables, duplicate variables, ...)" )
  
  def LOOP = "loop" ~> "{" ~> (BLOCK+) <~ "}" ^^ { case bs => Loop(bs) }
  def NDCHOICE = "choose" ~> "{" ~> (BLOCK+) ~ "}" ~ "{" ~ (BLOCK+) ~ "}" ^^  { case l ~ _ ~ _~ r ~ _ => NondetChoice(l,r) }
  
  def WORKFLOW = rep(BLOCK | LOOP) ^? { 
    case blocks if checkSig(blocks) => {
      val choiceblocks = addChoicePredicates(blocks)
      Workflow(buildSig(choiceblocks), choiceblocks) 
    }
  }
  
//  def WORKFLOW = BLOCK | (BLOCK ~ WORKFLOW) ^? { 
//    case block => { Workflow(Signature.EMPTY, List(block)) }
//    case block ~ workflow => { Workflow(Signature.EMPTY, workflow.steps) }
//  }
  
  def SPEC = 
    "Workflow" ~! WORKFLOW ~ 
    (("Declassify" ~> DECLASS)?) ~ 
    ("Target" ~> FUN) ~ 
    (("Causality" ~> repsep(TYPEDVAR, ","))?) ^^ { 
      case _ ~ w ~ decl ~ tar ~ causals => Spec(w, decl.toList, tar, causals.getOrElse(List())) 
    }
  
  def parseWorkflow(s: String) = parseAll(WORKFLOW, s)
  def parseSpec(s: String) = parseAll(SPEC, s)
  def parseTerm(s:String) = parseAll(TERM, s)
}