//! Ast definition for Ayin.

use std::collections::BTreeMap;

/// A Program.
#[derive(PartialEq, PartialOrd, Debug, Clone)]
pub struct Program {
    pub includes: Vec<File>,
    pub defs: Vec<Definition>,
}

#[derive(PartialEq, Eq, Hash, PartialOrd, Debug, Clone)]
pub struct File(pub String);

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

/// An expression.
#[derive(PartialEq, PartialOrd, Debug, Clone)]
pub enum Expr {
    Not(Box<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>),
}

#[derive(PartialEq, PartialOrd, Debug, Clone)]
pub enum Op {
    Assign,
    Compare(Cmp),
    Calc(Calc),
    Bool(BoolOp),
}

#[derive(PartialEq, PartialOrd, Debug, Clone)]
pub enum BoolOp {
    And,
    Or,
}

#[derive(PartialEq, PartialOrd, Debug, Clone)]
pub enum Cmp {
    Eq,
    NotEq,
    Gt,
    Lt,
    Gte,
    Lte,
}

#[derive(PartialEq, PartialOrd, Debug, Clone)]
pub enum Calc {
    Add,
    Sub,
    Div,
    Mul,
    Mod,
    BinAnd,
    BinOr,
}

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

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

/// 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, Hash, 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, Ref>);

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

// ----- //
// 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 Expr {
    fn from(rec: Vec<(&str, Expr)>) -> Expr {
        Expr::Record(
            rec.into_iter()
                .map(|(lbl, expr)| (Label(lbl.into()), expr))
                .collect(),
        )
    }
}

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<i32> for Value {
    fn from(i: i32) -> Value {
        Value::Int(i)
    }
}

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

impl From<f32> for Value {
    fn from(f: f32) -> Value {
        Value::Float(f)
    }
}

impl From<f32> for Expr {
    fn from(f: f32) -> Expr {
        Expr::Value(f.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,
            includes: vec![],
        }
    }
}

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

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

impl Vector {
    pub fn get(&self, index: usize) -> Result<&Ref, Error> {
        match self.0.get(index) {
            Some(v) => Ok(v),
            None => Err(Error::IndexOutOfBounds(index)),
        }
    }
}

impl std::fmt::Display for Expr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut string = "".to_string();
        match self {
            Expr::Value(value) => match value {
                Value::Int(i) => string.push_str(&format!("{i}")),
                Value::Float(f) => string.push_str(&format!("{f:2}")),
                Value::String(s) => string.push_str(&format!("{s:?}")),
                Value::Boolean(b) => string.push_str(&format!("{b}")),
                Value::Closure { .. } => string.push_str("<fn>"),
                _ => {}
            },
            Expr::Func(_) => string.push_str("<fn>"),
            Expr::Record(map) => {
                string.push_str("{");
                for (key, val) in map {
                    string.push_str(&format!(" .{}: {},", key.0, val));
                }
                if map.len() > 0 {
                    string.push_str(" }");
                } else {
                    string.push_str("}");
                }
            }
            Expr::Vector(vec) => {
                string.push_str("[");
                for val in vec {
                    string.push_str(&format!(" {},", val));
                }
                if vec.len() > 0 {
                    string.push_str(" ]");
                } else {
                    string.push_str("]");
                }
            }
            _ => {}
        }
        f.write_str(&string)
    }
}

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

pub enum Error {
    LabelNotFound(Label),
    IndexOutOfBounds(usize),
}

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Error::LabelNotFound(a) => {
                write!(f, "Label not found: {a:?}")
            }
            Error::IndexOutOfBounds(a) => {
                write!(f, "Indexout of bounds: {a:?}")
            }
        }
    }
}