//! Ast definition for Ayin.

use std::collections::BTreeMap;

/// An expression.
#[derive(PartialEq, PartialOrd, Debug, Clone)]
pub enum Expr {
    Op {
        lhs: Box<Expr>,
        op: Op,
        rhs: Box<Expr>,
    },
    Value(Value),
    Var(Name),
    Func(Box<Fn>),
    FunCall {
        func: Box<Expr>,
        args: Vec<Expr>,
    },
    Block(Vec<Statement>),
    If {
        condition: Box<Expr>,
        then: Box<Expr>,
        r#else: Box<Expr>,
    },
    Access {
        expr: Box<Expr>,
        field: Label,
    },
    Record(BTreeMap<Label, Expr>),
    Vector(Vec<Expr>),
    /*
    Loop {
        body: Box<Expr>,
    },
    Continue,
    Break,
    */
}

#[derive(PartialEq, PartialOrd, Debug, Clone)]
pub enum Op {
    Assign,
    //	Equals,
}

/// A statement.
#[derive(PartialEq, PartialOrd, Debug, Clone)]
pub enum Statement {
    Expr(Expr),
    Let(Definition),
    Return(Option<Expr>),
}

#[derive(PartialEq, PartialOrd, Debug, Clone)]
pub struct Definition {
    pub mutable: bool,
    pub name: Name,
    pub expr: Expr,
}

/// A reduced value.
#[derive(PartialEq, PartialOrd, Debug, Clone)]
pub enum Value {
    Int(i64),
    Float(f64),
    String(String),
    Boolean(bool),
    Record(Record),
    Vector(Vector),
    Label(Label),
    Ref(Ref),
    Closure { env: EnvName, expr: Box<Expr> },
}

/// A mutable variable.
#[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Clone)]
pub struct Ref(pub u64);

/// An anonymous function.
#[derive(PartialEq, PartialOrd, Debug, Clone)]
pub struct Fn {
    pub args: Vec<Arg>,
    pub body: Expr,
}

/// A function argument.
#[derive(PartialEq, PartialOrd, Debug, Clone)]
pub struct Arg {
    pub name: Name,
}

/// A symbol.
#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone)]
pub struct Name(pub String);

/// A field.
#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone)]
pub struct Label(pub String);

/// A reference to the environment.
#[derive(PartialEq, Eq, PartialOrd, Ord, Debug, Clone)]
pub struct EnvName(pub String);

#[derive(PartialEq, PartialOrd, Debug, Clone)]
pub struct Record(pub BTreeMap<Label, Value>);

#[derive(PartialEq, PartialOrd, Debug, Clone)]
pub struct Vector(pub Vec<Value>);

/// A Program.
#[derive(PartialEq, PartialOrd, Debug)]
pub struct Program(pub Vec<Definition>);

// ----- //
// Impls //
// ----- //

impl From<&str> for Name {
    fn from(value: &str) -> Name {
        Name(value.into())
    }
}

impl From<&str> for Label {
    fn from(value: &str) -> Label {
        Label(value.into())
    }
}

impl From<&str> for Expr {
    fn from(value: &str) -> Expr {
        Expr::Var(value.into())
    }
}

impl From<Vec<(&str, Expr)>> for Value {
    fn from(rec: Vec<(&str, Expr)>) -> Value {
        Value::Record(Record({
            rec.into_iter()
                .map(|(lbl, expr)| (Label(lbl.into()), expr))
                .collect()
        }))
    }
}

impl From<Vec<(&str, Expr)>> for Expr {
    fn from(rec: Vec<(&str, Expr)>) -> Expr {
        Expr::Value(rec.into())
    }
}

impl Expr {
    pub fn field(self, label: &str) -> Expr {
        Expr::Access {
            expr: Box::new(self),
            field: label.into(),
        }
    }
    pub fn call(self, args: Vec<Expr>) -> Expr {
        Expr::FunCall {
            func: Box::new(self),
            args,
        }
    }
}

impl From<i64> for Value {
    fn from(i: i64) -> Value {
        Value::Int(i)
    }
}

impl From<i64> for Expr {
    fn from(i: i64) -> Expr {
        Expr::Value(i.into())
    }
}

impl From<bool> for Value {
    fn from(b: bool) -> Value {
        Value::Boolean(b)
    }
}

impl From<bool> for Expr {
    fn from(b: bool) -> Expr {
        Expr::Value(b.into())
    }
}

impl From<Vec<Definition>> for Program {
    fn from(defs: Vec<Definition>) -> Program {
        Program(defs)
    }
}

impl From<Vec<Statement>> for Expr {
    fn from(stmts: Vec<Statement>) -> Expr {
        Expr::Block(stmts)
    }
}

impl Record {
    pub fn get(&self, label: &Label) -> Result<&Expr, Error> {
        match self.0.get(label) {
            Some(v) => Ok(v),
            None => Err(Error::LabelNotFound(label.clone())),
        }
    }
}

pub const UNIT: Expr = Expr::Value(Value::Record(Record(BTreeMap::new())));

#[derive(PartialEq, Debug, thiserror::Error)]
pub enum Error {
    #[error("Label not found: {0:?}")]
    LabelNotFound(Label),
}