/*
 * 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 typeck;
pub use typeck::Error as TypeckError;

use miette::SourceSpan;

#[derive(Clone, Debug)]
pub enum Pattern {
    Ident(Box<str>, parser::Offset),
    Record(Record<Pattern>),
    TyRecord(Record<Pattern>),
    Ignore,
}

impl core::cmp::PartialEq for Pattern {
    fn eq(&self, rhs: &Self) -> bool {
        match (self, rhs) {
            (Pattern::Ident(_, _), Pattern::Ident(_, _)) => true,
            (Pattern::Ignore, Pattern::Ignore) => true,
            (Pattern::Record(lxs), Pattern::Record(rxs)) => lxs == rxs,
            (Pattern::TyRecord(lxs), Pattern::TyRecord(rxs)) => lxs == rxs,
            // TODO: should Ident(_) be = Ignore?
            (_, _) => false,
        }
    }
    fn ne(&self, rhs: &Self) -> bool {
        match (self, rhs) {
            (Pattern::Ident(_, _), Pattern::Ident(_, _)) => false,
            (Pattern::Ignore, Pattern::Ignore) => false,
            (Pattern::Record(lxs), Pattern::Record(rxs)) => lxs != rxs,
            (Pattern::TyRecord(lxs), Pattern::TyRecord(rxs)) => lxs != rxs,
            (Pattern::Ident(_, _), Pattern::Ignore) | (Pattern::Ignore, Pattern::Ident(_, _)) => {
                false
            }
            (_, _) => true,
        }
    }
}

#[derive(Clone, Debug)]
pub struct Record<V> {
    pub span: SourceSpan,
    pub fields: Vec<(Option<Box<str>>, V)>,
}

impl<V: core::cmp::PartialEq> core::cmp::PartialEq for Record<V> {
    fn eq(&self, rhs: &Self) -> bool {
        self.fields == rhs.fields
    }
    fn ne(&self, rhs: &Self) -> bool {
        self.fields != rhs.fields
    }
}

#[derive(Clone, Debug)]
pub struct FullPattern {
    pub pat: Pattern,
    pub pty: Option<Box<(SourceSpan, Expr)>>,
}

impl core::cmp::PartialEq for FullPattern {
    fn eq(&self, rhs: &Self) -> bool {
        if self.pat != rhs.pat {
            return false;
        }
        match (&self.pty, &rhs.pty) {
            (None, None) => true,
            (Some(_), None) | (None, Some(_)) => false,
            (Some(lt), Some(rt)) => lt.1 == rt.1,
        }
    }
    fn ne(&self, rhs: &Self) -> bool {
        if self.pat == rhs.pat {
            return false;
        }
        match (&self.pty, &rhs.pty) {
            (None, None) => false,
            (Some(_), None) | (None, Some(_)) => true,
            (Some(lt), Some(rt)) => lt.1 != rt.1,
        }
    }
}

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

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

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

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

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

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

    Ref(usize),

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

    Select {
        prim_span: SourceSpan,
        prim: Box<Expr>,
        then: Vec<(SourceSpan, Box<str>)>,
    },
}

impl Pattern {
    pub fn push_to_penv(&self, env: &mut ParseEnv<'_>) {
        match self {
            Pattern::Ident(i, _) => env.names.push(i.clone()),
            Pattern::Record(xs) | Pattern::TyRecord(xs) => {
                xs.fields[..].iter().for_each(|i| i.1.push_to_penv(env));
            }
            Pattern::Ignore => {}
        }
    }

    pub fn pop_from_penv(&self, env: &mut ParseEnv<'_>) {
        match self {
            Pattern::Ident(..) => {
                env.names.pop();
            }
            Pattern::Record(xs) | Pattern::TyRecord(xs) => {
                xs.fields[..].iter().for_each(|i| i.1.pop_from_penv(env));
            }
            Pattern::Ignore => {}
        }
    }

    pub fn extract_exports(&self, outp: &mut Vec<Box<str>>) -> Result<(), Perr> {
        match self {
            Pattern::Ident(i, loc) => {
                let i2 = i.clone();
                if outp.iter().find(|j| i == *j).is_some() {
                    Err(Perr {
                        offset: usize::try_from(*loc).unwrap(),
                        kind: Pek::PatDupIdent(i2),
                    })
                } else {
                    outp.push(i2);
                    Ok(())
                }
            }
            Pattern::Record(xs) | Pattern::TyRecord(xs) => xs.fields[..]
                .iter()
                .try_for_each(|i| i.1.extract_exports(outp)),
            Pattern::Ignore => Ok(()),
        }
    }

    pub fn in_this_penv<T>(
        &self,
        env: &mut ParseEnv<'_>,
        f: impl FnOnce(&mut ParseEnv<'_>) -> T,
    ) -> T {
        let orig_height = env.names.len();
        self.push_to_penv(env);
        let height = env.names.len();
        let ret = f(env);
        assert_eq!(env.names.len(), height);
        self.pop_from_penv(env);
        assert_eq!(env.names.len(), orig_height);
        ret
    }
}

impl MaybeParse for Pattern {
    const DFL_CTX: &'static str = "pattern";

    fn maybe_parse(env: &mut ParseEnv<'_>) -> Result<Option<Self>, Perr> {
        let backup = env.lxr.clone();
        let Token {
            kind,
            offset: tok_start,
        } = match env.lxr.next().transpose()? {
            None => return Ok(None),
            Some(x) => x,
        };
        Ok(Some(match kind {
            Tok::PatOut(i) => Pattern::Ident(i, tok_start),
            Tok::PatIgnore => Pattern::Ignore,
            Tok::Caret | Tok::Dot => {
                let is_tylvl = kind == Tok::Caret;
                if let Some(r) = Record::<Pattern>::maybe_parse(env)? {
                    if is_tylvl {
                        Pattern::TyRecord(r)
                    } else {
                        Pattern::Record(r)
                    }
                } else {
                    let Token { kind, offset } = env.lxr.next_in_noeof("[^.]pattern")?;
                    return Err(parser::unexpected_token(offset, kind, "[^.]pattern"));
                }
            }
            _ => {
                env.lxr = backup;
                return Ok(None);
            }
        }))
    }
}

impl MaybeParse for FullPattern {
    const DFL_CTX: &'static str = "pattern";

    fn maybe_parse(env: &mut ParseEnv<'_>) -> Result<Option<Self>, Perr> {
        let pat = match Pattern::maybe_parse(env)? {
            None => return Ok(None),
            Some(x) => x,
        };

        // check for name collisions
        {
            let mut exports = Vec::new();
            pat.extract_exports(&mut exports)?;
        }

        let pty = if env.lxr.got(Tok::DubColon).is_some() {
            let start_loc = usize::try_from(env.lxr.offset()).unwrap();
            let ptyx = Expr::parse(env)?;
            let end_loc = usize::try_from(env.lxr.offset()).unwrap();
            Some(Box::new(((start_loc..end_loc).into(), ptyx)))
        } else {
            None
        };

        Ok(Some(FullPattern { pat, pty }))
    }
}

impl<T: Parse> MaybeParse for Record<T> {
    const DFL_CTX: &'static str = "record";

    fn maybe_parse(env: &mut ParseEnv<'_>) -> Result<Option<Self>, Perr> {
        let start_offset = match env.lxr.got(Tok::LBrace) {
            None => return Ok(None),
            Some(x) => x,
        };

        let mut fields: Vec<(Option<Box<str>>, _)> = Vec::new();

        loop {
            let lxrbak = env.lxr.clone();
            let name = {
                let Token { kind, offset } = env.lxr.next_in_noeof("record")?;

                if let Tok::DotIdent(i) = kind {
                    if env.lxr.expect(Tok::Assign, "record =").is_ok() {
                        if fields
                            .iter()
                            .find(|(j, _)| j.as_ref() == Some(&i))
                            .is_some()
                        {
                            env.lxr = lxrbak;
                            return Err(Perr {
                                offset: usize::try_from(offset).unwrap(),
                                kind: Pek::RcdDupIdent(i),
                            });
                        }
                        Some(i)
                    } else {
                        None
                    }
                } else if let Tok::RBrace = kind {
                    env.lxr = lxrbak;
                    break;
                } else {
                    None
                }
            };
            if name.is_none() {
                // backtrack
                env.lxr = lxrbak;
            }

            let expr = T::parse(env)?;
            env.lxr.expect(Tok::SemiColon, "record ;")?;
            fields.push((name, expr));
        }

        let end_offset = env.lxr.expect(Tok::RBrace, "record")?;
        Ok(Some(Record {
            span: (usize::try_from(start_offset).unwrap()..usize::try_from(end_offset).unwrap())
                .into(),
            fields,
        }))
    }
}

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: SourceSpan = (usize::try_from(offset).unwrap()..ioffset_end).into();
                if let Expr::Select { then, .. } = &mut ret {
                    then.push((span, i));
                } else {
                    ret = Expr::Select {
                        prim_span: (usize::try_from(fi_offset).unwrap()..prim_end).into(),
                        prim: Box::new(ret),
                        then: vec![(span, i)],
                    };
                }
            }
            _ => {
                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 = parse_minexpr(env);
                assert_eq!(env.names.len(), reside_knamcnt);
                args.push(match arg {
                    Ok(x) => 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
    }
}