use peg::parser;

use crate::ast::*;

parser! {
    pub grammar parser() for str {

        // comment

        rule comment() = quiet!{ "//" [^'\n' |'\r']* ['\n' | '\r'] }

        // whitespace

        rule ws_char() = [' ' | '\t' | '\n' | '\r']
        rule skip() = ws_char() / comment()
        rule _() = quiet!{ skip()* }
        rule __() = quiet!{ skip()+ }

        rule ident_tail() = ['a'..='z' | 'A'..='Z' | '0'..='9' | '_' | '\'']
        rule wb() = !ident_tail()

        // keywords

        rule kw_let_set()       = "let set"       wb()
        rule kw_let_element()   = "let element"   wb()
        rule kw_let_signature() = "let signature" wb()
        rule kw_let_instance()  = "let instance"  wb()
        rule kw_record()        = "record"        wb()
        rule kw_variant()       = "variant"       wb()
        rule kw_theory()        = "theory"        wb()
        rule kw_case()          = "case"          wb()
        rule kw_of()            = "of"            wb()
        rule kw_for()           = "for"           wb()
        rule kw_set()           = "set"           wb()
        rule kw_Set()           = "Set"           wb()
        rule kw_Nat()           = "Nat"           wb()
        rule kw_Int()           = "Int"           wb()
        rule kw_Float()         = "Float"         wb()
        rule kw_Str()           = "Str"           wb()
        rule kw_Bool()          = "Bool"          wb()

        rule keyword() =
        kw_let_set() / kw_let_element() / kw_let_signature() / kw_let_instance()
        / kw_record()  / kw_variant()     / kw_theory()
        / kw_case()    / kw_of()          / kw_for()
        / kw_Set() / kw_set()
        / kw_Nat() / kw_Int() / kw_Float() / kw_Str() / kw_Bool()

        // identifiers

        rule lower_ident() -> String
        = !keyword() s:$(['a'..='z'] ident_tail()*) { s.to_string() }

        rule upper_ident() -> String
        = !keyword() s:$(['A'..='Z'] ident_tail()*) { s.to_string() }

        rule elem_var() -> String = lower_ident()
        rule inst_var() -> String = lower_ident()
        rule set_var()  -> String = upper_ident()
        rule sig_var()  -> String = upper_ident()

        // projections, injections
        // Note: we don't capture the dot

        rule project() -> String
        = __ "." n:$(['a'..='z' | 'A'..='Z'] ident_tail()*) { n.to_string() }

        rule inject() -> String
        = _ n:$(['a'..='z' | 'A'..='Z'] ident_tail()*) "." __ { n.to_string() }

        rule project_lower() -> String
        = __ "." n:$(['a'..='z'] ident_tail()*) { n.to_string() }

        rule inject_lower() -> String
        = _ n:$(['a'..='z'] ident_tail()*) "." __ { n.to_string() }

        rule project_upper() -> String
        = __ "." n:$(['A'..='Z'] ident_tail()*) { n.to_string() }


        // literals

        rule nat_lit() -> Literal
        = n:$(['0'..='9']+) !"." { Literal::Nat(n.parse().unwrap()) }

        rule int_lit() -> Literal
        = "-" n:$(['0'..='9']+) !"." { Literal::Int(-(n.parse::<i64>().unwrap())) }

        rule float_lit() -> Literal
        = s:$("-"? ['0'..='9']+ "." ['0'..='9']+) { Literal::Float(s.parse().unwrap()) }

        rule str_lit() -> Literal
        = "\"" s:$((!"\"" [_])*) "\"" { Literal::Str(s.to_string()) }

        rule bool_lit() -> Literal
        = "true" { Literal::Bool(true) } / "false" { Literal::Bool(false) }

        rule literal() -> Literal
        = float_lit() / int_lit() / nat_lit() / str_lit() / bool_lit()

        // built-in

        rule builtin() -> BuiltIn
        = kw_Nat()   { BuiltIn::Nat }
        / kw_Int()   { BuiltIn::Int }
        / kw_Float() { BuiltIn::Float }
        / kw_Str()   { BuiltIn::Str }
        / kw_Bool()  { BuiltIn::Bool }

        // set layer

        rule set_field() -> RecordField
        = n:project_lower() _ ":" _ s:set() { RecordField { name: n, set: s } }

        rule variant_field() -> VariantField
        = n:inject_lower() _ ":" _ s:set() _ { VariantField { name: n, set: s } }

        rule claimed_set() -> Instance
        = kw_set() "(" _ i:instance() _ ")" { i }

        rule set() -> Set
        = kw_record()  _ "{" fs:(set_field() ** ",") _ "}" { Set::Record(fs) }
        / kw_variant() _ "{" _ vs:(variant_field() ** "|") "}" { Set::Variant(vs) }
        / b:builtin() { Set::BuiltIn(b) }
        / i:claimed_set() { Set::ClaimedSet(i) }
        / v:set_var() { Set::Var(v) }
        / "(" _ s:set() _ ")" { s }

        // signature layer

        rule param() -> Param
        = "(" _ n:elem_var() _ ":" _ s:set() _ ")" { Param { name: n, set: s } }

        rule param_list() -> Vec<Param> = param() ++ _

        rule sig_field() -> SigField
        = n:project_upper() _ "::" _ s:signature() { SigField { name: n, signature: s } }

        rule signature() -> Signature
        = kw_Set()    { Signature::Set }
        / kw_theory() _ "{" fs:(sig_field() ** ",") _ "}" { Signature::Theory(fs) }
        / ps:param_list() _ "->" _ cod:signature() { Signature::Ext { params: ps, codomain: Box::new(cod) } }
        / v:sig_var() { Signature::Var(v) }
        / "(" _ s:signature() _ ")" { s }

        // element layer

        rule elem_assign() -> ElemAssign
        = n:project() _ "=" _ e:element() { ElemAssign { name: n, element: e } }

        rule atom_elem() -> Element
        = l:literal() { Element::Literal(l) }
        / v:elem_var() { Element::Var(v) }
        / "{" fs:(elem_assign() ** ",") _ "}" { Element::Record(fs) }
        / "(" _ e:element() _ ")" { e }

        rule dot_elem() -> Element
        = tags:(t:inject() _ { t })* head:atom_elem() projs:(p:project() { p })*
        {
            let base = projs.into_iter().fold(head, |acc, p| {
                Element::Project{element: Box::new(acc), field: p}}
            );
            tags.into_iter().rev().fold(base, |acc, t| {
                Element::Inject{field: t, element: Box::new(acc)}
            })
        }

        rule app_elem() -> Element
        = head:dot_elem() tail:(__ d:dot_elem() { d })*
        {
            tail.into_iter().fold(head, |acc, a| {
                Element::App(Box::new(acc), Box::new(a))
            })
        }

        rule case_arm() -> CaseArm
        = t:inject() _ x:elem_var() _ "=>" _ body:element() { CaseArm { tag: t, bound: x, body } }

        rule element() -> Element
        = kw_case() __ scrut:element() _ kw_of() _ "{" arms:(_ a:case_arm() _ { a }) ** "|"  _ "}" { Element::Case { scrutinee: Box::new(scrut), arms } }
        / app_elem()

        // instance layer

        rule inst_assign() -> InstAssign
        = n:project() _ "=" _ i:instance()
        { InstAssign { name: n, instance: i } }

        rule atom_inst() -> Instance
        = s:set() { Instance::SetCoerce(Box::new(s)) }
        / v:inst_var() { Instance::Var(v) }
        / "{" fs:(inst_assign() ** ",") _ "}" { Instance::Record(fs) }
        / "(" _ i:instance() _ ")" { i }

        rule dot_inst() -> Instance
        = head:atom_inst() projs:(p:project() { p })*
        {
            projs.into_iter().fold(head, |acc, p| {
                Instance::Project{instance: Box::new(acc), field: p}
            })
        }

        rule app_inst() -> Instance
        = head:dot_inst() tail:(__ a:dot_elem() { a })*
        {
            tail.into_iter().fold(head, |acc, a| {
                Instance::App(Box::new(acc), Box::new(a))
            })
        }

        rule instance() -> Instance
        = kw_for() __ ps:param_list() _ "," _ body:instance() { Instance::For { params: ps, body: Box::new(body) } }
        / app_inst()

        // declarations

        rule decl() -> Decl
        = kw_let_set() __ n:set_var() _ "=" _ s:set() { Decl::Set { name: n, set: s } }
        / kw_let_element() __ n:elem_var() _ ":" _ s:set() _ "=" _ e:element() { Decl::Element { name: n, set: s, element: e } }
        / kw_let_signature() __ n:sig_var() _ "=" _ sg:signature() { Decl::Signature { name: n, signature: sg } }
        / kw_let_instance() __ n:inst_var() _ "::" _ sg:signature() _ "=" _ i:instance() { Decl::Instance { name: n, signature: sg, instance: i } }

        pub rule program() -> Programme
        = _ ds:(decl() ** _) _ { Programme(ds) }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn debug_parse(src: &str) {
        match parser::program(src) {
            Ok(p) => println!("```{}\n```\n=>\n{}\n", src, p),
            Err(e) => {
                let line = e.location.line;
                let col = e.location.column;
                let off = e.location.offset;
                println!("FAIL at {}:{} (offset {})", line, col, off);
                println!("expected: {:#}", e.expected);

                let before = &src[off.saturating_sub(40)..off];
                let after = &src[off..(off + 40).min(src.len())];
                println!("...{}⟨HERE⟩{}...", before, after);
                panic!("Parse failed!");
            }
        }
    }

    #[test]
    fn test_sets() {
        let src = r#"
let set Config = record {
    .enabled : Bool,
    .count   : Nat,
    .offset  : Int,
    .scale   : Float,
    .label   : Str
}

let set Maybe = variant {
    none. : record {}
  | some. : Config
}
        "#;

        debug_parse(src);
    }

    #[test]
    fn test_elements() {
        let src = r#"
            let element foo : Nat =
                case some. config .count of {
                    none. ignore => 0
                  | some. n => n
                }
        "#;

        debug_parse(src);
    }

    #[test]
    fn test_theories_and_instances() {
        let src = r#"
            let signature Graph = theory {
                .Node :: Set,
                .Edge :: (s : Node) (t : Node) -> Set
            }

            let instance natPoset :: Graph = {
                .Node = Nat,
                .Edge = for (s : Nat) (t : Nat), Bool
            }

            let element node : set(natPoset .Node) = 7
        "#;

        debug_parse(src);
    }
}