hamelin_lib 0.9.2

use std::fmt;
use std::fmt::{Display, Formatter};

use crate::write_utils::Indent;
use std::str::FromStr;

use crate::antlr::hamelinparser::{
    IdentifierContextAll, IdentifierContextAttrs, SimpleIdentifierContextAll,
};
use crate::antlr::{parse_identifier, parse_simple_identifier};
use crate::err::{TranslationError, TranslationErrors};
use antlr_rust::tree::ParseTree;
use derive_more::{From, TryUnwrap};
use std::rc::Rc;

#[derive(Clone, Debug, PartialEq, Hash, Eq, From, TryUnwrap)]
pub enum Identifier {
    Simple(SimpleIdentifier),
    Compound(CompoundIdentifier),
}

impl Identifier {
    pub fn to_hamelin(&self) -> String {
        match self {
            Identifier::Simple(simple) => simple.to_hamelin(),
            Identifier::Compound(compound) => compound.to_hamelin(),
        }
    }

    pub fn simples(&self) -> Vec<SimpleIdentifier> {
        match self {
            Identifier::Simple(simple) => vec![simple.clone()],
            Identifier::Compound(compound) => compound.as_slice().to_vec(),
        }
    }

    pub fn prefixed_with(&self, prefix: &SimpleIdentifier) -> CompoundIdentifier {
        let rest = self.simples();
        CompoundIdentifier::new(prefix.clone(), rest[0].clone(), rest[1..].to_vec())
    }

    pub fn from_simples(simples: Vec<SimpleIdentifier>) -> Option<Identifier> {
        match simples.as_slice() {
            [] => None,
            [single] => Some(single.clone().into()),
            [first, second, rest @ ..] => {
                Some(CompoundIdentifier::new(first.clone(), second.clone(), rest.to_vec()).into())
            }
        }
    }

    pub fn prefixed_with_str(&self, prefix: &str) -> CompoundIdentifier {
        self.prefixed_with(&SimpleIdentifier::new(prefix))
    }

    pub fn prefix(&self) -> &[SimpleIdentifier] {
        match self {
            Identifier::Simple(_) => &[],
            Identifier::Compound(compound) => {
                let s = compound.as_slice();
                &s[..s.len() - 1]
            }
        }
    }

    pub fn last(&self) -> &SimpleIdentifier {
        match self {
            Identifier::Simple(simple) => simple,
            Identifier::Compound(compound) => {
                let s = compound.as_slice();
                &s[s.len() - 1]
            }
        }
    }

    pub fn first(&self) -> &SimpleIdentifier {
        match self {
            Identifier::Simple(simple) => simple,
            Identifier::Compound(compound) => compound.first(),
        }
    }

    pub fn rest(&self) -> &[SimpleIdentifier] {
        match self {
            Identifier::Simple(_) => &[],
            Identifier::Compound(compound) => &compound.as_slice()[1..],
        }
    }

    pub fn fmt_indented(
        &self,
        f: &mut std::fmt::Formatter<'_>,
        indentation: Indent,
    ) -> std::fmt::Result {
        match self {
            Identifier::Simple(simple) => simple.fmt_indented(f, indentation),
            Identifier::Compound(compound) => compound.fmt_indented(f, indentation),
        }
    }

    pub fn prepend(self, id: Identifier) -> CompoundIdentifier {
        CompoundIdentifier::from_two_idents(id, self.into())
    }
}

impl Display for Identifier {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}

impl FromStr for Identifier {
    type Err = TranslationErrors;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        HamelinIdentifier::new(parse_identifier(s.to_string())?).to_sql()
    }
}

impl From<crate::tree::ast::identifier::Identifier> for Identifier {
    fn from(ast_ident: crate::tree::ast::identifier::Identifier) -> Self {
        match ast_ident {
            crate::tree::ast::identifier::Identifier::Simple(simple) => {
                Identifier::Simple(simple.into())
            }
            crate::tree::ast::identifier::Identifier::Compound(compound) => {
                Identifier::Compound(compound.into())
            }
        }
    }
}

impl From<crate::tree::ast::identifier::SimpleIdentifier> for Identifier {
    fn from(ast_ident: crate::tree::ast::identifier::SimpleIdentifier) -> Self {
        Identifier::Simple(ast_ident.into())
    }
}

impl From<crate::tree::ast::identifier::CompoundIdentifier> for Identifier {
    fn from(ast_ident: crate::tree::ast::identifier::CompoundIdentifier) -> Self {
        Identifier::Compound(ast_ident.into())
    }
}

#[derive(Clone, Debug, PartialEq, Hash, Eq)]
pub struct SimpleIdentifier {
    pub name: String,
}

impl SimpleIdentifier {
    pub fn new(name: &str) -> Self {
        Self {
            name: name.to_string(),
        }
    }

    pub fn from_parse_tree(tree: &SimpleIdentifierContextAll) -> Result<Self, TranslationError> {
        match tree {
            SimpleIdentifierContextAll::BackQuotedIdentifierContext(ctx) => {
                let raw = ctx.get_text();
                let text = raw
                    .get(1..raw.len().saturating_sub(1))
                    .ok_or_else(|| TranslationError::msg(ctx, "invalid back-quoted identifier"))?
                    .replace("``", "`");
                Ok(Self::new(&text))
            }
            SimpleIdentifierContextAll::UnquotedIdentifierContext(ctx) => {
                Ok(Self::new(&ctx.get_text()))
            }
            SimpleIdentifierContextAll::Error(ctx) => {
                Err(TranslationError::msg(ctx, "parse_error"))
            }
        }
    }

    pub fn is_clean(&self) -> bool {
        match parse_simple_identifier(self.name.clone()) {
            Ok(tree) => match tree.as_ref() {
                SimpleIdentifierContextAll::BackQuotedIdentifierContext(_) => true,
                SimpleIdentifierContextAll::UnquotedIdentifierContext(_) => true,
                SimpleIdentifierContextAll::Error(_) => false,
            },
            Err(_) => false,
        }
    }

    pub fn to_hamelin(&self) -> String {
        if self.is_clean() {
            self.name.clone()
        } else {
            format!("`{}`", self.name.replace("`", "``"))
        }
    }

    pub fn fmt_indented(&self, f: &mut Formatter<'_>, _indentation: Indent) -> std::fmt::Result {
        write!(f, "\"{}\"", self.name.replace("\"", "\"\""))
    }
}

// Conversions from tree AST identifiers to SQL identifiers
impl From<crate::tree::ast::identifier::SimpleIdentifier> for SimpleIdentifier {
    fn from(ast_ident: crate::tree::ast::identifier::SimpleIdentifier) -> Self {
        SimpleIdentifier::new(ast_ident.as_str())
    }
}

impl Display for SimpleIdentifier {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}

impl From<SimpleIdentifier> for String {
    fn from(ident: SimpleIdentifier) -> Self {
        ident.name
    }
}

impl FromStr for SimpleIdentifier {
    type Err = TranslationErrors;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        HamelinSimpleIdentifier::new(parse_simple_identifier(s.to_string())?).to_sql()
    }
}

#[derive(Clone, Debug, PartialEq, Hash, Eq)]
pub struct CompoundIdentifier {
    simples: Vec<SimpleIdentifier>,
}

impl CompoundIdentifier {
    pub fn new(
        first: SimpleIdentifier,
        second: SimpleIdentifier,
        rest: Vec<SimpleIdentifier>,
    ) -> Self {
        let mut simples = vec![first, second];
        simples.extend(rest);
        Self { simples }
    }

    pub fn as_slice(&self) -> &[SimpleIdentifier] {
        &self.simples
    }

    pub fn into_vec(self) -> Vec<SimpleIdentifier> {
        self.simples
    }

    pub fn to_hamelin(&self) -> String {
        self.simples
            .iter()
            .map(|s| s.to_hamelin())
            .collect::<Vec<String>>()
            .join(".")
    }

    pub fn from_two_idents(first: Identifier, second: Identifier) -> Self {
        Self::from_idents(first, second, &[])
    }

    pub fn from_idents(first: Identifier, second: Identifier, rest: &[Identifier]) -> Self {
        let all: Vec<_> = [first]
            .iter()
            .chain([second].iter())
            .chain(rest.iter())
            .flat_map(|id| match id {
                Identifier::Simple(simple) => vec![simple.clone()],
                Identifier::Compound(compound) => compound.as_slice().to_vec(),
            })
            .collect();
        match all.as_slice() {
            [first, second, rest @ ..] => Self::new(first.clone(), second.clone(), rest.to_vec()),
            _ => unreachable!("There are always at least two identifiers passed to from_idents"),
        }
    }

    pub fn from_two_str(one: &str, two: &str) -> Self {
        Self::new(
            SimpleIdentifier::new(one),
            SimpleIdentifier::new(two),
            vec![],
        )
    }

    pub fn first(&self) -> &SimpleIdentifier {
        &self.simples[0]
    }

    pub fn rest(&self) -> Identifier {
        if self.simples.len() == 2 {
            self.simples[1].clone().into()
        } else {
            CompoundIdentifier::new(
                self.simples[1].clone(),
                self.simples[2].clone(),
                self.simples[3..].to_vec(),
            )
            .into()
        }
    }

    pub fn fmt_indented(&self, f: &mut Formatter<'_>, indentation: Indent) -> std::fmt::Result {
        for (i, simple) in self.simples.iter().enumerate() {
            simple.fmt_indented(f, indentation)?;
            if i < self.simples.len() - 1 {
                write!(f, ".")?;
            }
        }
        Ok(())
    }
}

impl Display for CompoundIdentifier {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.fmt_indented(f, Indent::default())
    }
}

impl From<crate::tree::ast::identifier::CompoundIdentifier> for CompoundIdentifier {
    fn from(ast_ident: crate::tree::ast::identifier::CompoundIdentifier) -> Self {
        let all: Vec<SimpleIdentifier> = ast_ident
            .into_parts()
            .into_iter()
            .map(|p| p.into())
            .collect();
        CompoundIdentifier::new(all[0].clone(), all[1].clone(), all[2..].to_vec())
    }
}

// Need to keep this part of the "old AST" around until we figure out how to kill the SQL identifiers from the Type tree.
pub struct HamelinIdentifier {
    pub ctx: Rc<IdentifierContextAll<'static>>,
}

impl HamelinIdentifier {
    pub fn new(ctx: Rc<IdentifierContextAll<'static>>) -> Self {
        Self { ctx }
    }

    pub fn to_sql(&self) -> Result<Identifier, TranslationErrors> {
        let idents = TranslationErrors::from_vec(
            self.ctx
                .simpleIdentifier_all()
                .iter()
                .map(|si| HamelinSimpleIdentifier::new(si.clone()).to_sql())
                .collect(),
        )?;
        if idents.len() == 1 {
            Ok(idents[0].clone().into())
        } else if idents.is_empty() {
            Err(TranslationError::msg(self.ctx.as_ref(), "expected identifier").single())
        } else {
            match idents.as_slice() {
                [first, second, rest @ ..] => {
                    Ok(
                        CompoundIdentifier::new(first.clone(), second.clone(), rest.to_vec())
                            .into(),
                    )
                }
                _ => Err(TranslationError::msg(
                    self.ctx.as_ref(),
                    "expected at least two identifier parts",
                )
                .single()),
            }
        }
    }
}

pub struct HamelinSimpleIdentifier {
    pub ctx: Rc<SimpleIdentifierContextAll<'static>>,
}

impl HamelinSimpleIdentifier {
    pub fn new(ctx: Rc<SimpleIdentifierContextAll<'static>>) -> Self {
        Self { ctx }
    }

    pub fn to_sql(&self) -> Result<SimpleIdentifier, TranslationErrors> {
        SimpleIdentifier::from_parse_tree(self.ctx.as_ref()).map_err(|x| x.into())
    }

    pub fn get_tuple_index(&self) -> Result<usize, TranslationError> {
        let ctx = self.ctx.as_ref();
        let text = ctx.get_text();
        let msg = "Bad tuple reference. Must start with f and be followed by a number.";
        if !text.starts_with("f") {
            return Err(TranslationError::msg(ctx, msg));
        }
        text.get(1..)
            .ok_or_else(|| TranslationError::msg(ctx, msg))
            .and_then(|t| {
                t.parse::<usize>()
                    .map_err(|_| TranslationError::msg(ctx, msg))
            })
    }
}

#[cfg(test)]
mod test {
    use rstest::rstest;

    use super::*;

    #[rstest]
    #[case("identifier", "identifier")]
    #[case("`let`", "let")]
    #[case("`let.foob`", "let.foob")]
    pub fn test_simple_identifiers(#[case] hamelin: String, #[case] expected: String) {
        // First test that they parse as "identifiers"
        match HamelinIdentifier::new(parse_identifier(hamelin.clone()).unwrap())
            .to_sql()
            .unwrap()
        {
            Identifier::Simple(simple_identifier) => {
                assert_eq!(simple_identifier.name, expected)
            }
            Identifier::Compound(_) => panic!("should have been simple"),
        }

        // Then test that they parse as "simple identifiers"
        assert_eq!(
            HamelinSimpleIdentifier::new(parse_simple_identifier(hamelin).unwrap())
                .to_sql()
                .unwrap(),
            SimpleIdentifier::new(expected.as_str())
        );
    }

    #[rstest]
    #[case("one.two", vec!["one", "two"])]
    #[case("`let`.two", vec!["let", "two"])]
    #[case("`let`.`select`.`two.three`.four", vec!["let", "select", "two.three", "four"])]
    pub fn test_compound_identifiers(#[case] hamelin: String, #[case] expected: Vec<&str>) {
        // First test that they parse as "identifiers"
        match HamelinIdentifier::new(parse_identifier(hamelin.clone()).unwrap())
            .to_sql()
            .unwrap()
        {
            Identifier::Simple(_) => panic!("should have been compound"),
            Identifier::Compound(compound_identifier) => {
                assert_eq!(
                    compound_identifier
                        .into_vec()
                        .into_iter()
                        .map(|s| s.name)
                        .collect::<Vec<_>>(),
                    expected
                )
            }
        }
    }

    #[rstest]
    #[case("one.two")]
    #[case("select")]
    pub fn test_simple_identifier_errors(#[case] hamelin: String) {
        assert!(parse_simple_identifier(hamelin).is_err());
    }

    #[rstest]
    #[case("select")]
    pub fn test_identifier_errors(#[case] hamelin: String) {
        assert!(parse_simple_identifier(hamelin).is_err());
    }
}