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

use miette::SourceSpan;

use std::collections::VecDeque;

#[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, 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<(SourceSpan, Expr)>,
    },

    Ref(usize),

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

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

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,
        }))
    }
}

pub enum RecordEntry<'a> {
    Name(&'a str),
    Ref(usize),
}

impl<T> Record<T> {
    pub fn lookup_mut(&mut self, name: RecordEntry<'_>) -> Option<&mut T> {
        match name {
            RecordEntry::Name(name) => self
                .fields
                .iter_mut()
                .find(|(j, _)| j.as_ref().map(|x| &**x) == Some(name)),
            RecordEntry::Ref(i) => self.fields.get_mut(i),
        }
        .map(|x| &mut x.1)
    }

    pub fn lookup(&self, name: RecordEntry<'_>) -> Option<&T> {
        match name {
            RecordEntry::Name(name) => self
                .fields
                .iter()
                .find(|(j, _)| j.as_ref().map(|x| &**x) == Some(name)),
            RecordEntry::Ref(i) => self.fields.get(i),
        }
        .map(|x| &x.1)
    }
}

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_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
    }
}