/*
 * SPDX-FileCopyrightText: 2023 Alain Zscheile <fogti+devel@ytrizja.de>
 *
 * SPDX-License-Identifier: Apache-2.0
 */

//use bitflags::bitflags;
pub mod parser;
use parser::{
    Env as ParseEnv, Error as Perr, ErrorKind as Pek, MaybeParse, Parse, Token, TokenKind as Tok,
};

mod pat;
pub use pat::{FullPattern, Pattern};

mod record;
pub use record::{Record, RecordEntry};

pub mod typeck;

use core::cmp;
use miette::SourceSpan;
use std::collections::VecDeque;

/// A SourceSpan which is always equal to all other SourceSpans
/// (to be used to ignore span differences in expression comparisons)
#[derive(Clone, Copy, Debug)]
pub struct EvEqSourceSpan(pub SourceSpan);

impl cmp::PartialEq for EvEqSourceSpan {
    #[inline(always)]
    fn eq(&self, _: &Self) -> bool {
        true
    }
    #[inline(always)]
    fn ne(&self, _: &Self) -> bool {
        false
    }
}

impl From<(usize, usize)> for EvEqSourceSpan {
    #[inline]
    fn from(x: (usize, usize)) -> Self {
        Self(x.into())
    }
}

impl From<core::ops::Range<usize>> for EvEqSourceSpan {
    #[inline]
    fn from(x: core::ops::Range<usize>) -> Self {
        Self(x.into())
    }
}

impl From<EvEqSourceSpan> for SourceSpan {
    #[inline(always)]
    fn from(x: EvEqSourceSpan) -> SourceSpan {
        x.0
    }
}

impl<'a> From<&'a EvEqSourceSpan> for SourceSpan {
    #[inline(always)]
    fn from(x: &'a EvEqSourceSpan) -> SourceSpan {
        x.0
    }
}

pub trait Subst {
    fn incr_refs(&mut self, keep: usize, offset: usize);
    fn subst(&mut self, bnest: usize, with_e: &Expr);
}

#[derive(Clone, Debug, PartialEq)]
pub enum Expr {
    TyTy,
    Tbool,
    Tu32,
    Tusize,

    Lambda {
        pat: FullPattern,
        exp: Box<Expr>,
    },

    TyLambda {
        pat: FullPattern,
        body_span: EvEqSourceSpan,
        exp: Box<Expr>,
    },

    SelfRecur {
        pat: FullPattern,
        inner_span: EvEqSourceSpan,
        inner: Box<Expr>,
    },

    LetBind {
        kvs: Vec<(FullPattern, usize, Expr)>,
        inner: Box<Expr>,
    },

    Apply {
        lam: Box<Expr>,
        args: Vec<(EvEqSourceSpan, Expr)>,
    },

    Ref(usize),

    Record(Record<Expr>),
    TyRecord(Record<Expr>),

    Select {
        prim_span: EvEqSourceSpan,
        prim: Box<Expr>,
        then: VecDeque<(EvEqSourceSpan, Box<str>)>,
    },
}

fn parse_minexpr(env: &mut ParseEnv<'_>) -> Result<Expr, Perr> {
    let Token {
        offset: fi_offset,
        kind: fi_kind,
    } = env.lxr.next_in_noeof("expression")?;
    let mut ret = match fi_kind {
        Tok::Ident(i) => {
            if let Some(x) = env.lookup(&i) {
                Expr::Ref(x)
            } else {
                return Err(Perr {
                    offset: fi_offset.try_into().unwrap(),
                    kind: Pek::UnknownIdent(i),
                });
            }
        }
        Tok::Lambda => {
            let pat = FullPattern::parse(env)?;
            env.lxr.expect(Tok::RArr, "lambda")?;
            let exp = Box::new(pat.pat.in_this_penv(env, Expr::parse)?);
            return Ok(Expr::Lambda { pat, exp });
        }
        Tok::TyLambda => {
            let pat = FullPattern::parse(env)?;
            env.lxr.expect(Tok::RArr, "tylambda")?;
            let body_start = env.lxr.offset();
            let exp = Box::new(pat.pat.in_this_penv(env, Expr::parse)?);
            let body_end = env.lxr.offset();
            return Ok(Expr::TyLambda {
                pat,
                body_span: (body_start..body_end).into(),
                exp,
            });
        }
        Tok::Mu => {
            let pat = FullPattern::parse(env)?;
            env.lxr.expect(Tok::RArr, "mu")?;
            let inner_start = usize::try_from(env.lxr.offset()).unwrap();
            let inner = Box::new(pat.pat.in_this_penv(env, Expr::parse)?);
            let inner_end = usize::try_from(env.lxr.offset()).unwrap();
            return Ok(Expr::SelfRecur {
                pat,
                inner_span: (inner_start..inner_end).into(),
                inner,
            });
        }
        Tok::LParen => {
            let inner = Expr::parse(env)?;
            env.lxr.expect(Tok::RParen, "parens")?;
            return Ok(inner);
        }
        Tok::Caret => {
            return if let Some(r) = Record::maybe_parse(env)? {
                Ok(Expr::TyRecord(r))
            } else {
                let Token { kind, offset } = env.lxr.next_in_noeof("^expression")?;
                match kind {
                    Tok::Symbol(x) => match &*x {
                        "bool" => Ok(Expr::Tbool),
                        "type" => Ok(Expr::TyTy),
                        "u32" => Ok(Expr::Tu32),
                        "usize" => Ok(Expr::Tusize),
                        _ => Err(parser::unexpected_token(
                            offset,
                            Tok::Symbol(x),
                            "^expression",
                        )),
                    },
                    _ => Err(parser::unexpected_token(offset, kind, "^expression")),
                }
            }
        }
        Tok::Dot => {
            return if let Some(r) = Record::maybe_parse(env)? {
                Ok(Expr::Record(r))
            } else {
                let Token { kind, offset } = env.lxr.next_in_noeof(".expression")?;
                Err(parser::unexpected_token(offset, kind, ".expression"))
            }
        }
        /*
        Tok::LBrace => {
            let inner = ;
            env.lxr.expect(Tok::RBrace, "braces")?;
            return Ok(inner);
        }
        */
        _ => return Err(parser::unexpected_token(fi_offset, fi_kind, "expression")),
    };

    let prim_end = env.lxr.offset();

    loop {
        let lxrbak = env.lxr.clone();
        let Token { kind, offset } = match env.lxr.next() {
            Some(Ok(x)) => x,
            Some(Err(e)) => return Err(e),
            None => break,
        };
        match kind {
            Tok::DotIdent(i) => {
                let ioffset_end = env.lxr.offset();
                let span: EvEqSourceSpan = (usize::try_from(offset).unwrap()..ioffset_end).into();
                if let Expr::Select { then, .. } = &mut ret {
                    then.push_back((span, i));
                } else {
                    let mut then = VecDeque::new();
                    then.push_back((span, i));
                    ret = Expr::Select {
                        prim_span: (usize::try_from(fi_offset).unwrap()..prim_end).into(),
                        prim: Box::new(ret),
                        then,
                    };
                }
            }
            _ => {
                env.lxr = lxrbak;
                break;
            }
        }
    }
    Ok(ret)
}

impl Parse for Expr {
    fn parse(env: &mut ParseEnv<'_>) -> Result<Self, Perr> {
        fn parse_env_inner(env: &mut ParseEnv<'_>) -> Result<Expr, Perr> {
            let mut letbinds = Vec::new();

            while env.lxr.got(Tok::Let).is_some() {
                let key = FullPattern::parse(env)?;
                let middle = env.lxr.offset();
                env.lxr.expect(Tok::Assign, "let expression =")?;
                let value = Expr::parse(env)?;
                key.pat.push_to_penv(env);
                letbinds.push((key, middle, value));
                env.lxr.expect(Tok::SemiColon, "let expression ;")?;
            }

            let mut args = Vec::new();
            let mut base = parse_minexpr(env)?;
            let reside_knamcnt = env.names.len();
            loop {
                let lxrbak = env.lxr.clone();
                let arg_start = env.lxr.offset();
                let arg = parse_minexpr(env);
                let arg_end = env.lxr.offset();
                assert_eq!(env.names.len(), reside_knamcnt);
                match arg {
                    Ok(x) => args.push(((arg_start..arg_end).into(), x)),
                    Err(_) => {
                        // do not eat errors without backtracking the lexer
                        env.lxr = lxrbak;
                        break;
                    }
                }
            }

            if !args.is_empty() {
                base = Expr::Apply {
                    lam: Box::new(base),
                    args,
                };
            }

            if !letbinds.is_empty() {
                base = Expr::LetBind {
                    kvs: letbinds,
                    inner: Box::new(base),
                };
            }

            Ok(base)
        }

        // make sure we don't leak names
        // (the assert_eq checks serve the same purpose)
        let knamcnt = env.names.len();
        let res = parse_env_inner(env);
        assert!(env.names.len() >= knamcnt);
        env.names.truncate(knamcnt);
        res
    }
}

impl Expr {
    pub fn weak_is_type(&self) -> Option<bool> {
        Some(match self {
            Expr::TyTy | Expr::Tbool | Expr::Tu32 | Expr::Tusize => true,
            Expr::Lambda { .. } | Expr::Record(_) => false,
            Expr::TyLambda { .. } | Expr::TyRecord(_) => true,
            Expr::SelfRecur { inner, .. }
            | Expr::LetBind { inner, .. }
            | Expr::Apply { lam: inner, .. } => return inner.weak_is_type(),
            Expr::Ref(_) | Expr::Select { .. } => return None,
        })
    }
}

impl Subst for Expr {
    fn incr_refs(&mut self, mut keep: usize, offset: usize) {
        // this pattern prevents too deep recursion
        let mut this = Some(self);
        while let Some(cur) = this.take() {
            match cur {
                Expr::TyTy | Expr::Tbool | Expr::Tu32 | Expr::Tusize => {}

                Expr::Lambda { pat, exp }
                | Expr::TyLambda { pat, exp, .. }
                | Expr::SelfRecur {
                    pat, inner: exp, ..
                } => {
                    pat.incr_refs(keep, offset);
                    keep += pat.count_exports();
                    this = Some(exp);
                }

                Expr::LetBind { kvs, inner } => {
                    for (pat, _, value) in kvs {
                        pat.incr_refs(keep, offset);
                        value.incr_refs(keep, offset);
                        keep += pat.count_exports();
                    }
                    this = Some(inner);
                }

                Expr::Apply { lam, args } => {
                    lam.incr_refs(keep, offset);
                    for (_, i) in args {
                        i.incr_refs(keep, offset);
                    }
                }

                Expr::Ref(i) => {
                    if *i >= keep {
                        *i += offset;
                    }
                }

                Expr::Record(rcd) | Expr::TyRecord(rcd) => {
                    for (_, i) in &mut rcd.fields {
                        i.incr_refs(keep, offset);
                    }
                }

                Expr::Select { prim, .. } => this = Some(prim),
            }
        }
    }

    fn subst(&mut self, mut bnest: usize, with_e: &Expr) {
        match self {
            Expr::TyTy | Expr::Tbool | Expr::Tu32 | Expr::Tusize => {}

            Expr::Lambda { pat, exp }
            | Expr::TyLambda { pat, exp, .. }
            | Expr::SelfRecur {
                pat, inner: exp, ..
            } => {
                pat.subst(bnest, with_e);
                let mut with_e2 = with_e.clone();
                let patexps = pat.count_exports();
                with_e2.incr_refs(0, patexps);
                exp.subst(bnest + patexps, &with_e2);
            }

            Expr::LetBind { kvs, inner } => {
                let mut with_e2 = with_e.clone();
                for (pat, _, value) in kvs {
                    pat.subst(bnest, &with_e2);
                    value.subst(bnest, &with_e2);
                    let patexps = pat.count_exports();
                    with_e2.incr_refs(0, patexps);
                    bnest += patexps;
                }
                inner.subst(bnest, &with_e2);
            }

            Expr::Apply { lam, args } => {
                lam.subst(bnest, with_e);
                for (_, i) in args {
                    i.subst(bnest, with_e);
                }
            }

            Expr::Ref(i) => {
                if *i == bnest {
                    *self = with_e.clone();
                }
            }

            Expr::Record(rcd) | Expr::TyRecord(rcd) => {
                for (_, i) in &mut rcd.fields {
                    i.subst(bnest, with_e);
                }
            }

            Expr::Select { prim, .. } => prim.subst(bnest, with_e),
        }
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn subst00() {
        use crate::*;
        let mut a = Expr::Lambda {
            pat: FullPattern {
                pat: Pattern::Ident("x".to_string().into_boxed_str(), 0),
                pty: None,
            },
            exp: Box::new(Expr::Ref(0)),
        };
        a.subst(0, &Expr::Ref(1));
        assert_eq!(&a, &Expr::Lambda {
            pat: FullPattern {
                pat: Pattern::Ident("x".to_string().into_boxed_str(), 0),
                pty: None,
            },
            exp: Box::new(Expr::Ref(0)),
        });
        if let Expr::Lambda { exp, .. } = &mut a {
            **exp = Expr::Ref(1);
        }
        a.subst(0, &Expr::Ref(1));
        assert_eq!(&a, &Expr::Lambda {
            pat: FullPattern {
                pat: Pattern::Ident("x".to_string().into_boxed_str(), 0),
                pty: None,
            },
            exp: Box::new(Expr::Ref(2)),
        });
    }
}