use peg::parser;

use crate::ast::*;

parser! {
    grammar parser() for str {

        // ====================================================================
        // Whitespace and comments
        // ====================================================================

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

        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_element_of()    = "element-of"    wb()
        rule kw_for()           = "for"           wb()
        rule kw_set()           = "set"           wb()
        rule kw_set_of()        = "set-of"        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 kw_true()          = "true"          wb()
        rule kw_false()         = "false"         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_of()      / kw_set()           / kw_element_of()
        / kw_Nat()    / kw_Int()         / kw_Float()         / kw_Str()          / kw_Bool()
        / kw_true()   / kw_false()

        // ====================================================================
        // 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 any_ident() -> String
        = !keyword() s:$(['_' | 'a'..='z' | '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 (.field) and injections (tag.) -- chain forms
        // ====================================================================

        rule project() -> String
        = __ "." n:any_ident() { n }

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

        // ====================================================================
        // Record-field labels
        // ====================================================================

        rule field_lower() -> String = "." n:lower_ident() { n }
        rule field_upper() -> String = "." n:upper_ident() { n }

        // ====================================================================
        // Literals and built-in sets
        // ====================================================================

        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
        = kw_true()  { Literal::Bool(true) }
        / kw_false() { Literal::Bool(false) }

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

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

        // ====================================================================
        // Sets
        // ====================================================================

        rule set_field() -> Field<Set>
        = _ n:lower_ident() _ ":" _ s:set() _
        { Field { name: n, carries: s } }

        rule variant_field() -> Field<Set>
        = _ n:lower_ident() _ ":" _ s:set() _
        { Field { name: n, carries: s } }

        rule claimed_set() -> Instance
        = kw_set_of() _ "(" _ 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 }

        // ====================================================================
        // Signatures
        // ====================================================================

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

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

        rule sig_field() -> Field<Signature>
        = _ n:upper_ident() _ "::" _ s:signature() _
        { Field { name: n, carries: s } }

        rule sig_ext() -> Signature
        = ps:param_list() _ "->" _ cod:signature()
        {
            match cod {
                Signature::Ext { params: inner, codomain } => {
                    let mut all = ps;
                    all.extend(inner);
                    Signature::Ext { params: all, codomain }
                }
                other => Signature::Ext { params: ps, codomain: Box::new(other) },
            }
        }

        rule signature() -> Signature
        = kw_Set()                                           { Signature::Set }
        / "<" _ s:set() _ ">"                                { Signature::FromSet(s) }
        / kw_theory() _ "{" _ fs:(sig_field() ** ",") _ "}"  { Signature::Theory(fs) }
        / s:sig_ext()                                        { s }
        / v:sig_var()                                        { Signature::Var(v) }
        / "(" _ s:signature() _ ")"                          { s }

        // ====================================================================
        // Elements
        // ====================================================================

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

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

        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 case_arm() -> CaseArm<Element>
        = _ t:inject() _ x:elem_var() _ "=>" _ body:element() _
        { CaseArm { tag: t, bound: x, body } }

        rule claimed_element() -> Instance
        = kw_element_of() _ "(" _ i:instance() _ ")" { i }


        rule element() -> Element
        = i:claimed_element()                                                          { Element::ClaimedElement(Box::new(i)) }
        / kw_case() _ scrut:element() _ kw_of() _ "[" _ arms:(case_arm() ** "|") _ "]" { Element::Case { scrutinee: Box::new(scrut), arms } }
        / d:dot_elem()                                                                 { d }

        // ====================================================================
        // Instances
        // ====================================================================

        rule inst_case_arm() -> CaseArm<Instance>
        = _ t:inject() _ x:elem_var() _ "=>" _ body:instance() _
        { CaseArm { tag: t, bound: x, body } }

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

        rule explicit_set_coerce() -> Set
        = s:set() _ "::" _ kw_Set()
        { s }

        rule atom_inst() -> Instance
        = s:explicit_set_coerce()                    { Instance::SetCoerce(Box::new(s)) }
        / "<" _ e:element() _ ">"                    { Instance::ElementCoerce(e) }
        / v:inst_var()                               { Instance::Var(v) }
        / f:sig_var()                                { Instance::Var(f) }
        / "{" _ 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
        = subject:dot_inst() args:(__ a:dot_elem() { a })*
        {
            if args.is_empty() {
                subject
            } else {
                match subject {
                    Instance::App { instance: inner_head, args: inner_args } => {
                        let mut all_args = inner_args;
                        all_args.extend(args);
                        Instance::App { instance: inner_head, args: all_args }
                    }
                    other => Instance::App { instance: Box::new(other), args },
                }
            }
        }

        rule inst_for() -> Instance
        = kw_for() _ ps:param_list() _ "," _ body:instance()
        {
            match body {
                Instance::For { params: inner, body } => {
                    let mut all = ps;
                    all.extend(inner);
                    Instance::For { params: all, body }
                }
                other => Instance::For { params: ps, body: Box::new(other) },
            }
        }

        rule instance() -> Instance
        = i:inst_for()                                                                      { i }
        / kw_case() _ scrut:element() _ kw_of() _ "[" _ arms:(inst_case_arm() ** "|") _ "]" { Instance::Case { scrutinee: Box::new(scrut), arms } }
        / a:app_inst()                                                                      { a }

        // ====================================================================
        // Top-level 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) }
    }
}

pub fn parse_result(src: &str) -> Result<Programme, String> {
    match parser::program(src) {
        Ok(p) => Ok(p),
        Err(e) => {
            let line = e.location.line;
            let col = e.location.column;
            let off = e.location.offset;

            let before = &src[off.saturating_sub(40)..off];
            let after = &src[off..(off + 40).min(src.len())];

            Err(format!(
                "FAIL at {}:{} (offset {})\nexpected: {}\n...{}⟨HERE⟩{}...",
                line, col, off, e.expected, before, after
            ))
        }
    }
}