// SPDX-License-Identifier: MIT

use std::collections::HashSet;

use super::*;
use itertools::Itertools;
use krust_sitter_types::grammar::Grammar;
use syn::{parse::Parse, punctuated::Punctuated, spanned::Spanned};

#[derive(Debug)]
pub struct RuleDerive {
    pub ident: syn::Ident,
    pub attrs: Vec<Attribute>,
    pub params: RuleParams,
    pub data: syn::Data,
}

impl RuleDerive {
    pub fn from_derive_input(d: DeriveInput) -> Result<Option<Self>> {
        if d.attrs.iter().any(|a| {
            let Ok(list) = a.meta.require_list() else {
                return false;
            };
            let Some(derives) = list
                .parse_args_with(Punctuated::<Path, Token![,]>::parse_terminated)
                .ok()
            else {
                // Should be infallible to parse derive tokens
                return false;
            };
            derives
                .iter()
                .any(|p| p == &parse_quote!(krust_sitter::Rule) || p == &parse_quote!(Rule))
        }) {
            Ok(Some(Self::from_derive_input_known(d)?))
        } else {
            Ok(None)
        }
    }

    // Used by the proc macro directly.
    pub fn from_derive_input_known(d: DeriveInput) -> Result<Self> {
        let extras = RuleParams::new(&d.attrs)?;
        Ok(Self {
            ident: d.ident,
            attrs: d.attrs,
            params: extras,
            data: d.data,
        })
    }
}

/// Generate a single grammar per module.
pub fn generate_grammar(root_file: Vec<Item>) -> Result<Option<Grammar>> {
    let mut state = ExpansionState::new();
    // for some reason, source_file must be the first key for things to work
    state
        .grammar
        .rules
        .insert("source_file".to_string(), RuleDef::BLANK);

    if root_file
        .into_iter()
        .map(|item| process_item(item, &mut state))
        .flat_map(|r| r.err())
        .reduce(|mut acc, e| {
            acc.combine(e);
            acc
        })
        .map(Err::<(), _>)
        .is_some()
    {
        state.err()?;
        // Theoretically, we can use this for something. But, this generate_grammar function is
        // only used in `build.rs`, and all of these errors which are useful will be caught
        // during macro expansion, which is much more useful for development purposes.
        return Ok(None);
    }
    if state.verify_seen().is_err() {
        return Ok(None);
    }

    // This error is useful for us and cannot be generated by proc macro expansion.
    let (language_ident, language) = state.language_rule.ok_or_else(|| {
        Error::new(
            Span::call_site(),
            "Must specify exactly one root with #[language]",
        )
    })?;
    let name = language
        .name()
        .unwrap_or_else(|| language_ident.to_string());
    state.grammar.name = name.clone();
    state.grammar.rules.insert(
        "source_file".to_string(),
        state
            .grammar
            .rules
            .get(&language_ident.to_string())
            .unwrap()
            .clone(),
    );
    Ok(Some(state.grammar))
}

pub struct ExpansionState {
    pub grammar: Grammar,
    pub language_rule: Option<(Ident, LanguageExpr)>,
    // Accumulated errors.
    pub error: Option<Error>,
}

impl Default for ExpansionState {
    fn default() -> Self {
        Self::new()
    }
}

impl ExpansionState {
    pub fn new() -> Self {
        Self {
            grammar: Grammar {
                name: String::new(),
                word: None,
                rules: Default::default(),
                extras: Default::default(),
                conflicts: Default::default(),
            },
            language_rule: None,
            error: None,
        }
    }
    fn err(&mut self) -> Result<()> {
        if let Some(err) = self.error.take() {
            Err(err)
        } else {
            Ok(())
        }
    }

    fn accumulate_error(&mut self, err: Error) -> Error {
        if let Some(inner) = &mut self.error {
            inner.combine(err.clone());
        } else {
            self.error = Some(err.clone());
        }
        err
    }

    fn verify_seen(&self) -> Result<()> {
        if let Some(e) = self
            .grammar
            .rules
            .values()
            .flat_map(|v| self.check_seen_value(v).err())
            .reduce(|mut acc, e| {
                acc.combine(e);
                acc
            })
        {
            Err(e)
        } else {
            Ok(())
        }
    }

    fn check_seen_value(&self, value: &RuleDef) -> Result<()> {
        // Each value is always a map.
        match value {
            RuleDef::SYMBOL { name } => {
                if !self.grammar.rules.contains_key(name) {
                    return Err(Error::new(
                        Span::call_site(),
                        format!("Symbol found with no corresponding value: {name}"),
                    ));
                }
            }
            RuleDef::CHOICE { members } | RuleDef::SEQ { members } => {
                for member in members {
                    self.check_seen_value(member)?;
                }
            }
            RuleDef::REPEAT { content }
            | RuleDef::REPEAT1 { content }
            | RuleDef::PREC_DYNAMIC { value: _, content }
            | RuleDef::PREC_LEFT { value: _, content }
            | RuleDef::PREC_RIGHT { value: _, content }
            | RuleDef::PREC { value: _, content }
            | RuleDef::TOKEN { content }
            | RuleDef::IMMEDIATE_TOKEN { content }
            | RuleDef::RESERVED {
                context_name: _,
                content,
            } => self.check_seen_value(content)?,
            _ => return Ok(()),
        }
        Ok(())
    }

    fn set_language(&mut self, ident: &Ident, lang: LanguageExpr) -> Result<()> {
        if let Some((existing, _)) = &self.language_rule {
            return Err(self.accumulate_error(Error::new(
                lang.path.span(),
                format!(
                    "Language rule already defined as {existing}, found duplicate with {ident}",
                ),
            )));
        }
        self.language_rule = Some((ident.clone(), lang));
        Ok(())
    }

    fn set_word(&mut self, ident: String) -> Result<()> {
        if let Some(existing) = &self.grammar.word {
            return Err(self.accumulate_error(Error::new(
                Span::call_site(),
                format!("Word rule already defined as {existing}, found duplicate with {ident}"),
            )));
        }
        self.grammar.word = Some(ident);
        Ok(())
    }
}

fn process_item(item: Item, ctx: &mut ExpansionState) -> Result<()> {
    match item {
        Item::Struct(_) | Item::Enum(_) => {
            // Try and convert it to a derive.
            let stream = item.to_token_stream();
            let input = RuleDerive::from_derive_input(syn::parse2(stream)?)?;
            if let Some(input) = input {
                // Parse the structure now.
                process_rule(input, ctx)?;
            }
        }
        Item::Mod(m) => {
            // Recursively process this now.
            let (_, items) = m
                .content
                .ok_or_else(|| Error::new(Span::call_site(), "Module must be inlined"))?;
            for item in items {
                process_item(item, ctx)?;
            }
        }
        _ => {}
    }

    Ok(())
}

pub fn process_rule(input: RuleDerive, ctx: &mut ExpansionState) -> Result<()> {
    if let Some(l) = &input.params.language {
        ctx.set_language(&input.ident, l.clone())?;
    }
    if let Some(word) = &input.params.word {
        let rule = word.evaluate()?;
        match rule {
            RuleDef::SYMBOL { name } => ctx.set_word(name)?,

            _ => {
                return Err(Error::new(
                    word.span(),
                    "word must be a reference to another rule",
                ));
            }
        }
    }
    if let Some(extras) = &input.params.extras {
        if input.params.language.is_none() {
            return Err(Error::new(
                extras.span(),
                "Cannot specify extras without #[language]",
            ));
        }
        let (extras, errs): (Vec<_>, Vec<_>) = extras
            .iter()
            .map(|input| input.evaluate())
            .partition_result();
        let err = errs.into_iter().reduce(|mut acc, n| {
            acc.combine(n);
            acc
        });
        if let Some(err) = err {
            return Err(err);
        }
        ctx.grammar.extras = extras;
    }
    // if input.extras.word {
    //     ctx.set_word(&input.ident);
    // }

    let ident = input.ident;

    match input.data {
        Data::Struct(DataStruct { fields, .. }) => {
            gen_struct_or_variant(ident.to_string(), &input.attrs, fields.clone(), false, ctx)?;
        }
        Data::Enum(DataEnum { variants, .. }) => {
            variants
                .iter()
                .flat_map(|v| {
                    gen_struct_or_variant(
                        format!("{}_{}", ident, v.ident),
                        &v.attrs,
                        v.fields.clone(),
                        true,
                        ctx,
                    )
                    .err()
                })
                .reduce(|mut acc, e| {
                    acc.combine(e);
                    acc
                })
                .map(Err::<(), _>)
                .transpose()?;

            let mut members = vec![];
            variants.iter().for_each(|v| {
                let variant_path = format!("{}_{}", ident, v.ident);
                members.push(RuleDef::SYMBOL { name: variant_path });
            });

            let rule = RuleDef::CHOICE { members };

            let rule = input.params.apply(rule)?;

            ctx.grammar.rules.insert(ident.to_string(), rule);
        }
        Data::Union(_) => return Err(Error::new(ident.span(), "Union not supported")),
    }

    Ok(())
}

fn gen_field(
    ident_str: &str,
    leaf_type: Option<Type>,
    attrs: Vec<Attribute>,
) -> Result<(RuleDef, bool)> {
    let precs = RuleParams::new(&attrs)?;

    if precs.language.is_some() {
        return Err(Error::new(
            leaf_type.span(),
            "Cannot specify language on a field",
        ));
    }
    let leaf_attr = attrs.iter().find(|attr| sitter_attr_matches(attr, "leaf"));

    let text_attr = attrs.iter().find(|attr| sitter_attr_matches(attr, "text"));

    if let Some(leaf_attr) = leaf_attr
        && text_attr.is_some()
    {
        return Err(Error::new(
            leaf_attr.span(),
            "Cannot specify leaf and text at the same time",
        ));
    }

    if let Some(text) = text_attr {
        let input: TsInput = text.parse_args()?;
        return Ok((precs.apply(input.evaluate()?)?, false));
    }

    let leaf_input = leaf_attr.map(|a| a.parse_args::<TsInput>()).transpose()?;

    let leaf_type = match leaf_type {
        Some(ty) => ty,
        None => {
            let Some(leaf_input) = leaf_input else {
                // TODO: Narrow the span
                return Err(Error::new(
                    Span::call_site(),
                    "Empty types must have a leaf or text attribute",
                ));
            };
            return Ok((precs.apply(leaf_input.evaluate()?)?, false));
        }
    };

    let mut skip_over = HashSet::new();
    skip_over.insert("Spanned");
    skip_over.insert("Box");

    let (inner_type_vec, is_vec) = try_extract_inner_type(&leaf_type, "Vec", &skip_over);
    let (inner_type_option, is_option) = try_extract_inner_type(&leaf_type, "Option", &skip_over);

    if !is_vec && !is_option {
        if let Some(input) = leaf_input {
            let result = input.evaluate()?;
            Ok((precs.apply(result)?, is_option))
        } else {
            let symbol_name = match filter_inner_type(&leaf_type, &skip_over) {
                Type::Path(p) => p.path.require_ident()?.to_string(),
                t => return Err(Error::new(t.span(), "Expected a path")),
            };

            Ok((precs.apply(RuleDef::SYMBOL { name: symbol_name })?, false))
        }
    } else if is_vec {
        let (field_json, field_optional) = gen_field(
            ident_str,
            Some(inner_type_vec),
            leaf_attr.iter().cloned().cloned().collect(),
        )?;

        let (delimited_param, repeat_non_empty) = attrs
            .iter()
            .find_map(|attr| {
                if sitter_attr_matches(attr, "sep_by") {
                    Some((Some(attr), false))
                } else if sitter_attr_matches(attr, "sep_by1") {
                    Some((Some(attr), true))
                } else if sitter_attr_matches(attr, "repeat1") {
                    Some((None, true))
                } else {
                    None
                }
            })
            .unwrap_or((None, false));
        let delimited_param = delimited_param
            .map(|a| a.parse_args::<TsInput>())
            .transpose()?;

        let delimiter_json = delimited_param
            .as_ref()
            .map(|p| precs.apply(p.evaluate()?))
            .transpose()?;

        let field_rule_non_optional = field_json;
        let field_rule = if field_optional {
            RuleDef::optional(field_rule_non_optional)
        } else {
            field_rule_non_optional
        };

        let vec_contents = if let Some(delimiter_json) = delimiter_json {
            let content = Box::new(RuleDef::SEQ {
                members: vec![delimiter_json, field_rule.clone()],
            });
            let delimiter_rule = if field_optional {
                RuleDef::REPEAT1 { content }
            } else {
                RuleDef::REPEAT { content }
            };
            RuleDef::SEQ {
                members: vec![field_rule, delimiter_rule],
            }
        } else {
            RuleDef::REPEAT1 {
                // This FIELD is used only for the macro generation phase to distinguish between
                // the two different types of REPEAT1.
                content: Box::new(RuleDef::FIELD {
                    name: ident_str.to_owned(),
                    content: field_rule.into(),
                }),
            }
        };

        let vec_contents = precs.apply(vec_contents)?;

        Ok((vec_contents, !repeat_non_empty))
    } else {
        // is_option
        let (field_json, field_optional) = gen_field(ident_str, Some(inner_type_option), attrs)?;

        if field_optional {
            return Err(Error::new(
                Span::call_site(),
                "Option<Option<_>> is not supported",
            ));
        }

        Ok((precs.apply(field_json)?, true))
    }
}

fn gen_struct_or_variant(
    path: String,
    attrs: &[Attribute],
    fields: Fields,
    is_variant: bool,
    ctx: &mut ExpansionState,
) -> Result<()> {
    fn gen_field_optional(field: &Field, ident_str: String) -> Result<RuleDef> {
        let (field_contents, is_option) =
            gen_field(&ident_str, Some(field.ty.clone()), field.attrs.clone())?;

        let core = RuleDef::FIELD {
            name: ident_str,
            content: field_contents.into(),
        };

        let r = if is_option {
            RuleDef::optional(core)
        } else {
            core
        };
        Ok(r)
    }

    let (children, errs): (Vec<_>, Vec<_>) = fields
        .iter()
        .enumerate()
        .filter_map(|(i, field)| {
            if field
                .attrs
                .iter()
                .any(|attr| sitter_attr_matches(attr, "skip"))
            {
                None
            } else {
                let ident_str = field
                    .ident
                    .as_ref()
                    .map(|v| v.to_string())
                    .unwrap_or(format!("{i}"));

                Some(gen_field_optional(field, ident_str))
            }
        })
        .partition_result();
    let err = errs.into_iter().reduce(|mut acc, e| {
        acc.combine(e);
        acc
    });
    if let Some(err) = err {
        return Err(err);
    }

    let base_rule = match fields {
        Fields::Unit => {
            let (field_contents, _is_option) = gen_field("unit", None, attrs.to_owned())?;
            if is_variant {
                RuleDef::FIELD {
                    name: "unit".to_owned(),
                    content: field_contents.into(),
                }
            } else {
                field_contents
            }
        }
        _ => RuleDef::SEQ { members: children },
    };

    let precs = RuleParams::new(attrs)?;

    ctx.grammar.rules.insert(path, precs.apply(base_rule)?);
    Ok(())
}

#[derive(Debug)]
pub struct RuleParams {
    pub prec_param: Option<Expr>,
    pub prec_left_param: Option<Expr>,
    pub prec_right_param: Option<Expr>,
    pub prec_dynamic_param: Option<Expr>,
    pub language: Option<LanguageExpr>,
    pub extras: Option<Punctuated<TsInput, Token![,]>>,
    pub word: Option<TsInput>,
}

impl RuleParams {
    fn new(attrs: &[Attribute]) -> Result<Self> {
        let prec_attr = attrs.iter().find(|attr| sitter_attr_matches(attr, "prec"));

        let prec_param = prec_attr
            .map(|a| a.parse_args_with(Expr::parse))
            .transpose()?;

        let prec_left_attr = attrs
            .iter()
            .find(|attr| sitter_attr_matches(attr, "prec_left"));

        let prec_left_param = prec_left_attr
            .map(|a| a.parse_args_with(Expr::parse))
            .transpose()?;

        let prec_right_attr = attrs
            .iter()
            .find(|attr| sitter_attr_matches(attr, "prec_right"));

        let prec_right_param = prec_right_attr
            .map(|a| a.parse_args_with(Expr::parse))
            .transpose()?;

        let prec_dynamic_attr = attrs
            .iter()
            .find(|attr| sitter_attr_matches(attr, "prec_dynamic"));

        let prec_dynamic_param = prec_dynamic_attr
            .map(|a| a.parse_args_with(Expr::parse))
            .transpose()?;

        if let (Some(prec_left), Some(_prec_right)) = (prec_left_attr, prec_right_attr) {
            return Err(Error::new(
                prec_left.span(),
                "only one of prec, prec_left, and prec_right can be specified",
            ));
        }

        let extras = attrs
            .iter()
            .find(|a| sitter_attr_matches(a, "extras"))
            .map(|a| a.parse_args_with(Punctuated::<TsInput, Token![,]>::parse_terminated))
            .transpose()?;

        let language = attrs
            .iter()
            .find(|a| sitter_attr_matches(a, "language"))
            .map(LanguageExpr::from_attr)
            .transpose()?;
        let word = attrs
            .iter()
            .find(|a| sitter_attr_matches(a, "word"))
            .map(|a| a.parse_args_with(TsInput::parse))
            .transpose()?;
        // TODO: Refactor this and make it a sub-field of the `LanguageExpr` type instead.
        if language.is_none()
            && let Some(w) = &word
        {
            return Err(Error::new(
                w.span(),
                "Cannot specify #[word] on non-language rule",
            ));
        }

        Ok(Self {
            prec_param,
            prec_left_param,
            prec_right_param,
            prec_dynamic_param,
            extras,
            word,
            language,
        })
    }

    fn apply(&self, rule: RuleDef) -> Result<RuleDef> {
        let Self {
            prec_param,
            prec_left_param,
            prec_right_param,
            prec_dynamic_param,
            ..
        } = self;

        let rule = if let Some(Expr::Lit(lit)) = prec_param {
            if let Lit::Int(i) = &lit.lit {
                let value = i.base10_parse::<i32>()?;
                RuleDef::PREC {
                    value: value.into(),
                    content: Box::new(rule),
                }
            } else {
                return Err(Error::new(
                    lit.span(),
                    "Expected integer literal for precedence",
                ));
            }
        } else if let Some(Expr::Lit(lit)) = prec_left_param {
            let value = if let Lit::Int(i) = &lit.lit {
                i.base10_parse::<i32>()?
            } else {
                return Err(Error::new(
                    lit.span(),
                    "Expected integer literal for precedence",
                ));
            };
            RuleDef::PREC_LEFT {
                value: value.into(),
                content: Box::new(rule),
            }
        } else if let Some(Expr::Lit(lit)) = prec_right_param {
            let value = if let Lit::Int(i) = &lit.lit {
                i.base10_parse::<i32>()?
            } else {
                return Err(Error::new(
                    lit.span(),
                    "Expected integer literal for precedence",
                ));
            };
            RuleDef::PREC_RIGHT {
                value: value.into(),
                content: Box::new(rule),
            }
        } else if let Some(Expr::Lit(lit)) = prec_dynamic_param {
            if let Lit::Int(i) = &lit.lit {
                RuleDef::PREC_DYNAMIC {
                    value: i.base10_parse::<i32>()?,
                    content: Box::new(rule),
                }
            } else {
                return Err(Error::new(
                    lit.span(),
                    "Expected integer literal for precedence",
                ));
            }
        } else {
            rule
        };

        Ok(rule)
    }
}
