pub use super::common::Type as TypeNode;

use super::common::{BinaryOp, UnaryOp};
use super::location::*;

#[derive(Clone, Debug)]
pub struct Program {
  pub decls: Vec<Decl>,
  pub rules: Vec<Rule>,
  pub facts: Vec<Fact>,
  pub disjunctions: Vec<Disjunction>,
}

impl Node for TypeNode {
  type T = Self;

  fn new(t: Self) -> Self {
    t
  }
}

pub type Type = Located<TypeNode>;

#[derive(Clone, Debug)]
pub struct DeclNode {
  pub predicate: String,
  pub arg_types: Vec<Type>,
}

impl Node for DeclNode {
  type T = (String, Vec<Type>);

  fn new((predicate, arg_types): Self::T) -> Self {
    Self {
      predicate,
      arg_types,
    }
  }
}

pub type Decl = Located<DeclNode>;

#[derive(Clone, Debug)]
pub struct FactNode {
  pub prob: Option<f32>,
  pub head: Atom,
}

impl Node for FactNode {
  type T = (Option<f32>, Atom);

  fn new((prob, head): Self::T) -> Self {
    Self { prob, head }
  }
}

pub type Fact = Located<FactNode>;

#[derive(Clone, Debug)]
pub struct DisjunctionNode {
  pub facts: Vec<Fact>,
}

impl Node for DisjunctionNode {
  type T = Vec<Fact>;

  fn new(facts: Self::T) -> Self {
    Self { facts }
  }
}

pub type Disjunction = Located<DisjunctionNode>;

#[derive(Clone, Debug)]
pub struct RuleNode {
  pub head: Atom,
  pub body: Vec<Literal>,
}

impl Node for RuleNode {
  type T = (Atom, Vec<Literal>);

  fn new((head, body): Self::T) -> Self {
    Self { head, body }
  }
}

pub type Rule = Located<RuleNode>;

#[derive(Clone, Debug)]
pub struct AtomNode {
  pub predicate: String,
  pub args: Vec<Argument>,
}

impl Node for AtomNode {
  type T = (String, Vec<Argument>);

  fn new((predicate, args): Self::T) -> Self {
    Self { predicate, args }
  }
}

pub type Atom = Located<AtomNode>;

#[derive(Clone, Debug)]
pub enum LiteralNode {
  Pos(Atom),
  Neg(Atom),
  Constraint(Constraint),
}

impl Node for LiteralNode {
  type T = Self;

  fn new(data: Self) -> Self {
    data
  }
}

pub type Literal = Located<LiteralNode>;

#[derive(Clone, Debug)]
pub enum Constraint {
  Binary(BinaryConstraint),
  Unary(UnaryConstraint),
}

#[derive(Clone, Debug)]
pub struct BinaryConstraintNode {
  pub op: BinaryOp,
  pub op1: Argument,
  pub op2: Argument,
}

impl Node for BinaryConstraintNode {
  type T = (BinaryOp, Argument, Argument);

  fn new((op, op1, op2): Self::T) -> Self {
    Self { op, op1, op2 }
  }
}

pub type BinaryConstraint = Located<BinaryConstraintNode>;

#[derive(Clone, Debug)]
pub struct UnaryConstraintNode {
  pub op: UnaryOp,
  pub op1: Argument,
}

impl Node for UnaryConstraintNode {
  type T = (UnaryOp, Argument);

  fn new((op, op1): Self::T) -> Self {
    Self { op, op1 }
  }
}

pub type UnaryConstraint = Located<UnaryConstraintNode>;

#[derive(Clone, Debug)]
pub enum ConstantNode {
  Symbol(String),
  Boolean(bool),
  Integer(i64),
  SymbolId(usize),
  String(String),
}

impl Node for ConstantNode {
  type T = Self;

  fn new(t: Self::T) -> Self {
    t
  }
}

pub type Constant = Located<ConstantNode>;

#[derive(Clone, Debug)]
pub struct VariableNode {
  pub name: String,
}

impl Node for VariableNode {
  type T = String;

  fn new(name: Self::T) -> Self {
    Self { name }
  }
}

pub type Variable = Located<VariableNode>;

#[derive(Clone, Debug)]
pub enum Argument {
  Wildcard(Wildcard),
  Unary(UnaryExpr),
  Binary(BinaryExpr),
  Constant(Constant),
  Variable(Variable),
}

impl Argument {
  pub fn location(&self) -> &Location {
    match self {
      Self::Wildcard(w) => &w.location,
      Self::Unary(u) => &u.location,
      Self::Binary(b) => &b.location,
      Self::Constant(c) => &c.location,
      Self::Variable(v) => &v.location,
    }
  }
}

#[derive(Clone, Debug)]
pub struct WildcardNode;

impl Node for WildcardNode {
  type T = ();

  fn new((): Self::T) -> Self {
    Self
  }
}

pub type Wildcard = Located<WildcardNode>;

#[derive(Clone, Debug)]
pub struct UnaryExprNode {
  pub op: UnaryOp,
  pub op1: Box<Argument>,
}

impl Node for UnaryExprNode {
  type T = (UnaryOp, Box<Argument>);

  fn new((op, op1): Self::T) -> Self {
    Self { op, op1 }
  }
}

pub type UnaryExpr = Located<UnaryExprNode>;

#[derive(Clone, Debug)]
pub struct BinaryExprNode {
  pub op: BinaryOp,
  pub op1: Box<Argument>,
  pub op2: Box<Argument>,
}

impl Node for BinaryExprNode {
  type T = (BinaryOp, Box<Argument>, Box<Argument>);

  fn new((op, op1, op2): Self::T) -> Self {
    Self { op, op1, op2 }
  }
}

pub type BinaryExpr = Located<BinaryExprNode>;
